diff --git a/.gitignore b/.gitignore
index 974b630a8..ef5e183bd 100644
--- a/.gitignore
+++ b/.gitignore
@@ -51,7 +51,7 @@ htmlcov/
.tox/
.coverage
.coverage.*
-.cache
+.cache*
nosetests.xml
coverage.xml
*.cover
@@ -61,7 +61,7 @@ coverage.xml
# Translations
*.mo
*.pot
-
+deepmath_103k
# Django stuff:
*.log
local_settings.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/exp/cold_start/train_cold_start.py b/cookbook/exp/cold_start/train_cold_start.py
deleted file mode 100644
index da7149bba..000000000
--- a/cookbook/exp/cold_start/train_cold_start.py
+++ /dev/null
@@ -1,332 +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, SamplerBackend,
- IntentClassifier, HardFilter, RefuseFilter, DeadLoopFilter, TokenSoupFilter, MessageSanityFilter,
- SpecialCharsFilter, ModelFilter, DedupFilter,
- MessageNormalizer,
-)
-
-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/condenser/dataset.py b/cookbook/exp/condenser/dataset.py
deleted file mode 100644
index 32c30de4b..000000000
--- a/cookbook/exp/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/condenser/make_condenser_dataset.py b/cookbook/exp/condenser/make_condenser_dataset.py
deleted file mode 100644
index cf56a44e3..000000000
--- a/cookbook/exp/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/condenser/train_condenser_ddp.py b/cookbook/exp/condenser/train_condenser_ddp.py
deleted file mode 100644
index 997235781..000000000
--- a/cookbook/exp/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/condenser/untested/eval_condensed.py b/cookbook/exp/condenser/untested/eval_condensed.py
deleted file mode 100644
index 730aaf3a8..000000000
--- a/cookbook/exp/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/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/build_thinking_rag_index.py b/cookbook/exp/embedding/build_thinking_rag_index.py
deleted file mode 100644
index d228a597a..000000000
--- a/cookbook/exp/embedding/build_thinking_rag_index.py
+++ /dev/null
@@ -1,935 +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
-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_lora_transformers/step_4000',
-)
-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))
-
-# Hard-templated hints: the condenser SFT prior maps `Skill` to the legacy
-# `Use when: / numbered steps / Output:` skeleton on long inputs; embedding the
-# exact 4-line body template + explicit negative constraints is the only way to
-# override it deterministically across query and cot sides.
-RAG_QUERY_HINT = (
- 'Summarize this query for retrieval. '
- 'The body of ## Summary MUST follow this EXACT 4-line template โ '
- 'do NOT emit "Use when:", numbered procedure steps, or "Output:":\n'
- 'Topic: \n'
- 'Problem: \n'
- 'Skill: \n'
- 'Knowledge: \n'
- 'Then emit the mandatory ## More section as usual. '
- 'Topic must name the specific pattern, never generic labels.')
-RAG_THINKING_HINT = (
- 'Summarize this reasoning trace for retrieval. '
- 'The body of ## Summary MUST follow this EXACT 4-line template โ '
- 'do NOT emit "Use when:", numbered procedure steps, or "Output:":\n'
- 'Topic: \n'
- 'Problem: \n'
- 'Skill: \n'
- 'Knowledge: \n'
- 'Then emit the mandatory ## More section as usual. '
- 'Topic must name the specific pattern, never generic labels.')
-
-# OpenAI API fallback (used when vLLM truncates).
-COMPRESS_API_KEY = os.environ.get('COMPRESS_API_KEY', '')
-COMPRESS_BASE_URL = os.environ.get(
- '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)
-
-
-def _resolve_compressed(sampler: vLLMSampler, api: Optional[OpenAIClient],
- texts: List[str], query_hint: str) -> List[Optional[str]]:
- """Run vLLM batch; replace truncations / skeleton-incomplete with API output."""
- pairs = _vllm_compress(sampler, texts, query_hint)
- results: List[Optional[str]] = []
- for (text, stop), src_text in zip(pairs, texts):
- if stop != 'length' and not _is_truncated_compression(text):
- results.append(text)
- continue
- if api is None:
- results.append(None)
- continue
- api_text = _api_compress(api, _build_compress_messages(src_text, query_hint))
- if api_text is None or _is_truncated_compression(api_text):
- results.append(None)
- else:
- results.append(api_text)
- return results
-
-
-# ===========================================================================
-# Embedding helpers (TransformersModel.forward_only(task='embedding'))
-# ===========================================================================
-
-def _build_features(template: Qwen3_5Template, texts: List[str], role: str
- ) -> List[Dict[str, Any]]:
- """Wrap each text into the role-specific anchor / positive feature dict."""
- features: List[Dict[str, Any]] = []
- for text in texts:
- if not text or not text.strip():
- # Pad with a single space so positional alignment holds against
- # the input list โ the caller filters out empty-text rows upstream.
- text = ' '
- if role == 'anchor':
- feat = template.encode({'messages': _wrap_anchor(text)})
- feat['labels'] = [1]
- else:
- feat = template.encode({'messages': _wrap_positive(text)})
- feat['labels'] = [0]
- features.append(feat)
- return features
-
-
-def get_embeddings(model: TransformersModel, template: Qwen3_5Template,
- texts: List[str], role: str) -> np.ndarray:
- """Return ``[N, H]`` float32 L2-normalised embeddings for ``texts``.
-
- Inputs are padded up to a multiple of ``EMB_GPUS`` and sliced back to the
- original ``N``: the dispatch layer (``_dispatch_args``) starves any rank
- whose chunk lands beyond ``len(texts)``, so a single forward of fewer than
- ``EMB_GPUS`` items (e.g. the probe) would otherwise raise
- ``Batch too small for {EMB_GPUS} workers``.
- """
- if not texts:
- return np.zeros((0,), dtype=np.float32)
- n = len(texts)
- pad_n = (-n) % EMB_GPUS
- padded = list(texts) + [' '] * pad_n if pad_n else list(texts)
- features = _build_features(template, padded, role)
- out = model.forward_only(inputs=features, task='embedding', return_logits=True)
- emb = out['embeddings']
- if isinstance(emb, torch.Tensor):
- emb = emb.detach().to(torch.float32).cpu().numpy()
- emb = np.asarray(emb, dtype=np.float32)
- return emb[:n] if pad_n else emb
-
-
-def _probe_hidden_size(model: TransformersModel, template: Qwen3_5Template) -> int:
- """One-shot warmup forward to read out the embedding dimension."""
- emb = get_embeddings(model, template, ['probe'], role='anchor')
- if emb.ndim != 2 or emb.shape[0] == 0:
- raise RuntimeError(f'unexpected embedding shape from probe: {emb.shape}')
- return int(emb.shape[1])
-
-
-# ===========================================================================
-# LanceDB I/O
-# ===========================================================================
-
-def _make_arrow_schema(hidden_size: int):
- import pyarrow as pa
- return pa.schema([
- pa.field('id', pa.string()),
- pa.field('vector', pa.list_(pa.float32(), hidden_size)),
- pa.field('thinking_raw', pa.string()),
- pa.field('query_raw', pa.string()),
- pa.field('cot_compressed', pa.string()),
- pa.field('query_compressed', pa.string()),
- pa.field('source', pa.string()),
- pa.field('domain', pa.string()),
- pa.field('language', pa.string()),
- pa.field('sim', pa.float32()),
- ])
-
-
-def _open_or_create_table(db_path: str, table_name: str, hidden_size: int,
- mode: str):
- """Open an existing table for append/eval, or create a fresh one."""
- import lancedb
- db = lancedb.connect(db_path)
- schema = _make_arrow_schema(hidden_size)
- if table_name in db.table_names():
- if mode == 'overwrite':
- db.drop_table(table_name)
- tbl = db.create_table(table_name, schema=schema, mode='overwrite')
- else:
- tbl = db.open_table(table_name)
- else:
- tbl = db.create_table(table_name, schema=schema, mode='create')
- return db, tbl
-
-
-def _existing_ids(table) -> set:
- try:
- col = table.to_pandas(columns=['id'])
- return set(col['id'].astype(str).tolist())
- except Exception: # noqa: BLE001
- return set()
-
-
-# ===========================================================================
-# Build pipeline
-# ===========================================================================
-
-def _stream_corpus(total: Optional[int], load_from_cache_file: bool,
- max_rows: int = 0) -> Iterator[Dict[str, Any]]:
- ds = _GET_DATASET(total=total, load_from_cache_file=load_from_cache_file)
- n_full = len(ds)
- cap = max_rows if (max_rows and max_rows < n_full) else n_full
- sys.stderr.write(f'[corpus] get_dataset: {n_full} rows'
- + (f' โ yielding first {cap}\n' if cap < n_full else '\n'))
- for i, row in enumerate(ds):
- if i >= cap:
- break
- yield row
-
-
-def _extract_query_cot(row: Dict[str, Any]) -> Tuple[str, str]:
- user_query, cot = '', ''
- for m in row.get('messages') or []:
- if not isinstance(m, dict):
- continue
- role = m.get('role') or ''
- if role == 'user' and not user_query:
- user_query = (m.get('content') or '').strip()
- elif role == 'assistant':
- cot = (m.get('reasoning_content') or '').strip()
- break
- return user_query, cot
-
-
-def _log_miss(misses_path: str, lock: PosixFileLock, record: Dict[str, Any]) -> None:
- line = json.dumps(record, ensure_ascii=False, default=str) + '\n'
- with lock:
- with open(misses_path, 'a', encoding='utf-8') as fh:
- fh.write(line)
-
-
-def build_index(args: argparse.Namespace,
- sampler: vLLMSampler,
- emb_model: TransformersModel,
- emb_template: Qwen3_5Template,
- api: Optional[OpenAIClient]) -> None:
- # ---- Probe embedding dimension -----------------------------------------
- sys.stderr.write('[build] probing embedding hidden size...\n')
- hidden_size = _probe_hidden_size(emb_model, emb_template)
- sys.stderr.write(f'[build] hidden_size={hidden_size}\n')
-
- # ---- LanceDB ------------------------------------------------------------
- db, tbl = _open_or_create_table(
- args.db_path, args.table, hidden_size,
- mode='overwrite' if args.overwrite else 'append',
- )
- indexed = _existing_ids(tbl) if not args.overwrite else set()
- sys.stderr.write(f'[build] table "{args.table}" โ {len(indexed)} existing rows.\n')
-
- misses_path = args.misses_log or (str(Path(args.db_path) / f'{args.table}.misses.jsonl'))
- Path(misses_path).parent.mkdir(parents=True, exist_ok=True)
- misses_lock = PosixFileLock(misses_path + '.lock')
-
- # ---- Streaming loop -----------------------------------------------------
- n_seen = n_kept = n_dropped_short = n_dropped_compress = n_dropped_sim = 0
- n_dropped_dup = 0
- pbar = tqdm(desc='index', unit='row', dynamic_ncols=True)
-
- batch: List[Dict[str, Any]] = []
-
- def _flush(rows: List[Dict[str, Any]]) -> None:
- nonlocal n_kept, n_dropped_compress, n_dropped_sim
- if not rows:
- return
- # Phase 1 โ compress query (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)
- kept_rows: List[Dict[str, Any]] = []
- for r, q_cmp, c_cmp in zip(rows, q_compressed, c_compressed):
- if not q_cmp or not c_cmp:
- n_dropped_compress += 1
- _log_miss(misses_path, misses_lock, {
- 'id': r['id'], 'source': r['source'], 'reason': 'compress_fail',
- 'query_raw_head': _short(r['query_raw'], 200),
- 'cot_raw_head': _short(r['cot_raw'], 200),
- })
- continue
- r['query_compressed'] = q_cmp
- r['cot_compressed'] = c_cmp
- kept_rows.append(r)
- if not kept_rows:
- return
- # Phase 2 โ encode anchor (compressed query) + positive (compressed cot).
- anchor_emb = get_embeddings(
- emb_model, emb_template, [r['query_compressed'] for r in kept_rows], role='anchor')
- positive_emb = get_embeddings(
- emb_model, emb_template, [r['cot_compressed'] for r in kept_rows], role='positive')
- sims = (anchor_emb * positive_emb).sum(axis=1).astype(np.float32)
- # Phase 3 โ sim filter + LanceDB insert.
- to_insert: List[Dict[str, Any]] = []
- for idx, (r, sim_val) in enumerate(zip(kept_rows, sims)):
- tag = 'KEEP' if sim_val >= SIM_THRESHOLD else 'DROP'
- print(f'[{tag} sim={sim_val:.4f}] {r["source"][:24]} '
- f'q={_short(r["query_raw"], 60)!r} '
- f'cot={_short(r["cot_raw"], 60)!r}', flush=True)
- if sim_val < SIM_THRESHOLD:
- n_dropped_sim += 1
- _log_miss(misses_path, misses_lock, {
- 'id': r['id'], 'source': r['source'], 'reason': 'sim_low',
- 'sim': float(sim_val),
- 'query_raw': r['query_raw'],
- 'cot_raw': r['cot_raw'],
- 'query_compressed': r['query_compressed'],
- 'cot_compressed': r['cot_compressed'],
- })
- continue
- to_insert.append({
- 'id': r['id'],
- 'vector': positive_emb[idx].tolist(),
- 'thinking_raw': r['cot_raw'],
- 'query_raw': r['query_raw'],
- 'cot_compressed': r['cot_compressed'],
- 'query_compressed': r['query_compressed'],
- 'source': r['source'],
- 'domain': DOMAIN_MAP.get(r['source'], 'mixed'),
- 'language': _detect_lang(r['cot_raw']),
- 'sim': float(sim_val),
- })
- if to_insert:
- tbl.add(to_insert)
- n_kept += len(to_insert)
- indexed.update(r['id'] for r in to_insert)
-
- try:
- for row in _stream_corpus(total=args.total, load_from_cache_file=not args.no_cache,
- max_rows=args.max_rows):
- n_seen += 1
- if args.limit and n_kept >= args.limit:
- break
- rid = row.get('id') or ''
- if not rid:
- continue
- if rid in indexed:
- n_dropped_dup += 1
- continue
- user_query, cot = _extract_query_cot(row)
- if not user_query or len(cot) < MIN_TEXT_CHARS:
- n_dropped_short += 1
- continue
- batch.append({
- 'id': rid,
- 'source': row.get('source') or 'unknown',
- 'query_raw': user_query,
- 'cot_raw': cot,
- })
- if len(batch) >= args.batch_size:
- _flush(batch)
- batch.clear()
- pbar.set_postfix(kept=n_kept, sim_drop=n_dropped_sim,
- cmp_drop=n_dropped_compress, refresh=False)
- pbar.update(1)
- if batch:
- _flush(batch)
- batch.clear()
- finally:
- pbar.close()
-
- sys.stderr.write(
- f'[build] seen={n_seen} kept={n_kept} sim_drop={n_dropped_sim} '
- f'cmp_drop={n_dropped_compress} short_drop={n_dropped_short} '
- f'dup_skip={n_dropped_dup}\n')
-
- # ---- Build vector index for fast retrieval ------------------------------
- if n_kept >= 64 and not args.skip_index:
- sys.stderr.write('[build] creating IVF_PQ index (metric=dot)...\n')
- n_partitions = max(8, min(256, n_kept // 1000 + 1))
- try:
- tbl.create_index(
- metric='dot',
- vector_column_name='vector',
- num_partitions=n_partitions,
- num_sub_vectors=16,
- index_type='IVF_PQ',
- replace=True,
- )
- except Exception as exc: # noqa: BLE001
- sys.stderr.write(f'[build] index build failed: {exc} '
- '(table is still queryable via brute-force scan)\n')
- sys.stderr.write(f'[build] done. table rows={tbl.count_rows()}\n')
-
-
-# ===========================================================================
-# Eval pipeline (self-recall on indexed rows)
-# ===========================================================================
-
-def eval_recall(args: argparse.Namespace,
- sampler: vLLMSampler,
- emb_model: TransformersModel,
- emb_template: Qwen3_5Template,
- api: Optional[OpenAIClient]) -> None:
- """Probe each gold query against the index; report recall@k.
-
- Self-recall semantics: only rows whose ``id`` is already present in the
- index are probed. The corresponding ``cot``-keyed vector must be retrieved
- by encoding the **raw user query** through the condenser โ embedder
- pipeline (anchor side). The match is correct iff the retrieved row's
- ``id`` equals the probe row's ``id``.
- """
- import lancedb
- db = lancedb.connect(args.db_path)
- if args.table not in db.table_names():
- raise SystemExit(f'[eval] table "{args.table}" does not exist in {args.db_path}')
- tbl = db.open_table(args.table)
- indexed_ids = _existing_ids(tbl)
- sys.stderr.write(f'[eval] table rows={tbl.count_rows()} indexed_ids={len(indexed_ids)}\n')
- if not indexed_ids:
- sys.stderr.write('[eval] empty index โ nothing to evaluate.\n')
- return
-
- ks = sorted({1, 5, 10, args.top_k})
- hits = {k: 0 for k in ks}
- per_source_hits: Dict[str, Dict[int, int]] = {}
- per_source_total: Dict[str, int] = {}
- probed = 0
-
- pbar = tqdm(desc='eval', unit='probe', dynamic_ncols=True)
- batch_rows: List[Dict[str, Any]] = []
-
- def _flush(rows: List[Dict[str, Any]]) -> None:
- nonlocal probed
- if not rows:
- return
- compressed = _resolve_compressed(
- sampler, api, [r['query_raw'] for r in rows], 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='dataset_index',
- choices=['dataset_index', 'dataset_think'],
- help='Which loader to use: dataset_index (RAG profile) or '
- 'dataset_think (training mix).')
- p.add_argument('--limit', type=int, default=0,
- help='Stop building once this many rows are kept (0 = no cap).')
- p.add_argument('--max-rows', type=int, default=0,
- help='Truncate corpus to this many rows AFTER get_dataset (0 = no cap). '
- 'Use this instead of --total to avoid invalidating the dataset cache.')
- p.add_argument('--batch-size', type=int, default=64,
- help='Rows per condense+encode batch.')
- p.add_argument('--no-cache', action='store_true',
- help='Disable load_from_cache_file in dataset_think.get_dataset.')
- p.add_argument('--overwrite', action='store_true',
- help='Drop the table before build and start fresh.')
- p.add_argument('--skip-index', action='store_true',
- help='Skip IVF_PQ index build at the end (debug).')
- p.add_argument('--misses-log', default='',
- help='Path for filtered-row JSONL log (defaults to /.misses.jsonl).')
-
- # eval-only
- p.add_argument('--eval-size', type=int, default=500,
- help='Number of probes for self-recall evaluation.')
- p.add_argument('--top-k', type=int, default=10,
- help='Largest k to report. Smaller ks (1, 5) are always reported.')
-
- return p.parse_args()
-
-
-def main() -> None:
- args = parse_args()
- Path(args.db_path).mkdir(parents=True, exist_ok=True)
-
- global _GET_DATASET
- if args.dataset_module == 'dataset_think':
- from dataset_think import get_dataset as _swap
- _GET_DATASET = _swap
- sys.stderr.write(f'[main] dataset loader: {args.dataset_module}\n')
-
- # Build/eval both depend on the same Twinkle stack โ initialize once.
- sampler_mesh, emb_mesh = initialize_twinkle()
- sys.stderr.write(f'[main] twinkle initialized: '
- f'sampler ranks 0-{SAMPLER_GPUS - 1} (TP={SAMPLER_GPUS}), '
- f'emb_model ranks {SAMPLER_GPUS}-{NUM_GPUS - 1} (DP={EMB_GPUS}).\n')
-
- sys.stderr.write('[main] starting vLLM condenser sampler...\n')
- sampler = build_sampler(sampler_mesh)
- sys.stderr.write('[main] starting embedding TransformersModel...\n')
- emb_model, emb_template = build_emb_model(emb_mesh)
-
- api: Optional[OpenAIClient] = None
- if COMPRESS_API_KEY:
- api = OpenAIClient(
- model=COMPRESS_API_MODEL,
- api_key=COMPRESS_API_KEY,
- base_url=COMPRESS_BASE_URL,
- )
- else:
- sys.stderr.write(
- '[main] WARNING: COMPRESS_API_KEY unset โ truncated rows will be dropped.\n')
-
- if args.mode in ('build', 'both'):
- build_index(args, sampler, emb_model, emb_template, api)
- if args.mode in ('eval', 'both'):
- eval_recall(args, sampler, emb_model, emb_template, api)
-
-
-if __name__ == '__main__':
- main()
diff --git a/cookbook/exp/embedding/dataset_index.py b/cookbook/exp/embedding/dataset_index.py
deleted file mode 100644
index c86e1c523..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, batched=True, 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/train_embedding_full_ddp.py b/cookbook/exp/embedding/train_embedding_full_ddp.py
deleted file mode 100644
index bc69c56fd..000000000
--- a/cookbook/exp/embedding/train_embedding_full_ddp.py
+++ /dev/null
@@ -1,835 +0,0 @@
-"""LoRA embedding training with online compression via frozen vLLM condenser.
-
-Architecture (8 GPUs total):
- - Ranks 0-3 (``model``): Trainable embedding model with LoRA, InfoNCE loss.
- - Ranks 4-7 (``condenser_sampler``): Frozen vLLM condenser for online compression.
-
-When the condenser sampler truncates or regresses to the legacy schema, an
-external OpenAI-compatible API produces the correct compression. The failure is
-logged to failures.jsonl for offline SFT data regeneration.
-
-Launch:
- python cookbook/exp/train_embedding_lora_ddp.py
-"""
-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
-
-# -- 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))
-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
-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)
-
-
-# =============================================================================
-# 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)
-
-
-# =============================================================================
-# Compression prompt building
-# =============================================================================
-
-# Hard-templated hints: the condenser SFT prior maps `Skill` to the legacy
-# `Use when: / numbered steps / Output:` skeleton on long inputs; embedding the
-# exact 4-line body template + explicit negative constraints is the only way to
-# override it deterministically across query and cot sides.
-EMBED_QUERY_Q = (
- 'Summarize this query for retrieval. '
- 'The body of ## Summary MUST follow this EXACT 4-line template โ '
- 'do NOT emit "Use when:", numbered procedure steps, or "Output:":\n'
- 'Topic: \n'
- 'Problem: \n'
- 'Skill: \n'
- 'Knowledge: \n'
- 'Then emit the mandatory ## More section as usual. '
- 'Topic must name the specific pattern, never generic labels.')
-EMBED_QUERY_COT = (
- 'Summarize this reasoning trace for retrieval. '
- 'The body of ## Summary MUST follow this EXACT 4-line template โ '
- 'do NOT emit "Use when:", numbered procedure steps, or "Output:":\n'
- 'Topic: \n'
- 'Problem: \n'
- 'Skill: \n'
- 'Knowledge: \n'
- 'Then emit the mandatory ## More section as usual. '
- 'Topic must name the specific pattern, never generic labels.')
-
-# Legacy schema (Use when: / numbered steps / Output:) โ mixed in 50/50 with the
-# new schema to expose the embedder to schema-invariant semantic alignment.
-# Both query and cot of the SAME pair always use the SAME schema; cross-schema
-# anchors and positives would re-introduce the schema asymmetry we just fixed.
-EMBED_QUERY_Q_LEGACY = (
- 'What problem does this passage address, and what skill or method is needed? '
- 'Topic must name the specific pattern, never generic labels. '
- 'Compress into a retrieval-friendly need description.')
-EMBED_QUERY_COT_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
- if role == 'anchor':
- feat = template.encode({'messages': [
- {'role': 'user', 'content': decoded_text},
- {'role': 'assistant', 'content': 'Match the correct response here.'},
- ]})
- feat['labels'] = [1]
- else:
- feat = template.encode({'messages': [
- {'role': 'user', 'content': 'Match the correct query here.'},
- {'role': 'assistant', 'content': decoded_text},
- ]})
- 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
-
-
-# =============================================================================
-# Main training
-# =============================================================================
-
-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) ----------------------------------------
- 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)
-
- # -------- 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,
- )
-
- 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')
-
- 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,
- })
-
- # -------- 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
-
- # 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]
- 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:
- t0 = time.monotonic()
- minibatches = future.result() if future else None
- t_prefetch = time.monotonic() - t0
- future = prefetch_executor.submit(_sample_batch, raw_mega_batch)
-
- if not minibatches:
- continue
-
- for mb in minibatches:
- t1 = time.monotonic()
- model.forward_backward(inputs=mb, task='embedding')
- model.clip_grad_and_step(gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS)
- t_train = time.monotonic() - t1
- cur_step += 1
-
- _ema_prefetch = _ema_alpha * t_prefetch + (1 - _ema_alpha) * _ema_prefetch if cur_step > RESUME_STEP + 1 else t_prefetch
- _ema_train = _ema_alpha * t_train + (1 - _ema_alpha) * _ema_train if cur_step > RESUME_STEP + 1 else t_train
-
- if cur_step % LOG_INTERVAL == 0:
- metric = model.calculate_metric(is_training=True)
- _bottleneck = 'PREFETCH' if _ema_prefetch > _ema_train else 'TRAIN'
- logger.info(
- f'Epoch {epoch} Step {cur_step}/{total_forward_steps}, metric: {metric} | '
- f'prefetch={t_prefetch:.1f}s(ema {_ema_prefetch:.1f}) '
- f'train={t_train:.1f}s(ema {_ema_train:.1f}) '
- f'bottleneck={_bottleneck}')
- log_dict = {}
- for k, v in metric.items():
- if not v:
- continue
- try:
- log_dict[k] = float(v)
- except (ValueError, TypeError):
- pass
- log_dict['epoch'] = epoch
- log_dict['prefetch_sec'] = round(t_prefetch, 2)
- log_dict['train_sec'] = round(t_train, 2)
- swanlab.log(log_dict, step=cur_step)
- if cur_step % SAVE_INTERVAL == 0:
- save_checkpoint(model, f'step_{cur_step}')
- t_prefetch = 0.0
-
- # Drain final mega-batch
- if future:
- minibatches = future.result()
- future = None
- if minibatches:
- for mb in minibatches:
- model.forward_backward(inputs=mb, task='embedding')
- model.clip_grad_and_step(gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS)
- cur_step += 1
- if cur_step % SAVE_INTERVAL == 0:
- save_checkpoint(model, f'step_{cur_step}')
-
- prefetch_executor.shutdown(wait=False)
- save_checkpoint(model, 'last-checkpoint')
-
-
-if __name__ == '__main__':
- train()
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/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/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/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/grpo/kodcode_grpo.py b/cookbook/rl/grpo/kodcode_grpo.py
new file mode 100644
index 000000000..a7b44447a
--- /dev/null
+++ b/cookbook/rl/grpo/kodcode_grpo.py
@@ -0,0 +1,496 @@
+"""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
+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()
+
+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))
+
+_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>')
+
+
+# ========== Text handling (same as e18_kodcode) ==========
+def clean_text(decoded: Optional[str]) -> str:
+ return _SPECIAL_TOKEN_RE.sub('', decoded or '').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
+ jobs.append((i, unwrap_code(assistant_text(traj)), 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:
+ rewards.append(1.0 if unwrap_code(assistant_text(traj)).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..25673f563
--- /dev/null
+++ b/cookbook/rl/grpo/mbpp_grpo.py
@@ -0,0 +1,387 @@
+"""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 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
+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()
+
+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.')
+
+
+# ========== 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
+ jobs.append((i, unwrap_code(assistant_text(traj)), 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:
+ rewards.append(1.0 if unwrap_code(assistant_text(traj)).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/rsi/agentic/README.md b/cookbook/rsi/agentic/README.md
new file mode 100644
index 000000000..d1db30375
--- /dev/null
+++ b/cookbook/rsi/agentic/README.md
@@ -0,0 +1,285 @@
+# agentic โ RSI self-play where one trajectory is one request
+
+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.
+
+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.)
+
+## Three resources, three queues
+
+| 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:
+
+```
+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
+```
+
+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.
+
+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.
+
+`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.
+
+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
+sandbox.py the sandbox as a resource: clear, snapshot, run a script
+prompts.py every string sent to a model
+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
+rsi_agent.yaml the ms-agent config both sides' openings are shaped by
+```
+
+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
+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
+```
+
+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
+export E2B_API_KEY=... # sandbox host
+export SANDBOX_API_URL=http://... # sandbox host address, with port
+export LLM_BACKUP_API_KEY=... # dashscope
+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
+```
+
+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 `$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. 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
+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.
+
+## Settings that shape what gets produced
+
+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.
+
+| setting | value | origin |
+|---|---|---|
+| `--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: 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 |
+| `--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)` | 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
+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
new file mode 100644
index 000000000..b9ddaefc0
--- /dev/null
+++ b/cookbook/rsi/agentic/challenge.py
@@ -0,0 +1,1360 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""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, 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
+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 (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
+
+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.
+
+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
+import json
+import math
+import os
+import queue
+import random
+import statistics
+import sys
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple
+
+from twinkle import DeviceMesh, get_logger
+from twinkle.data_format import SamplingParams
+from twinkle.sampler import vLLMSampler
+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.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()
+
+# โโ 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 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
+
+
+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)
+ # 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',
+ 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=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)
+
+ # 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=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. 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)
+ p.add_argument('--propose-max-tokens', type=int, default=8192)
+ 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 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.
+ p.add_argument('--keywords-n', type=int, default=128,
+ 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)
+ p.add_argument('--keyword-max-tokens', type=int, default=4096)
+
+ # Novelty.
+ 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.")
+ # 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.')
+ # 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')
+ 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)
+
+ # Training, one step per iteration. See train.py.
+ 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',
+ 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 '
+ '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=16384,
+ help='a trajectory longer than this is not trained on. Below '
+ '--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')
+ 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-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.")
+
+ # 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.tag:
+ raise SystemExit('[rsi] --tag is required: it decides which run these '
+ 'iterations belong to and which checkpoint they overwrite')
+ # 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
+ return args
+
+
+# โโ Resources โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+
+def build_sampler(args) -> Tuple[Any, Any]:
+ """The resident vLLM sampler; returns (sampler, template).
+
+ 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.
+ """
+ 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(
+ '[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.')
+ 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,
+ 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,
+ )
+
+
+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
+
+
+# โโ 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.
+ """
+ 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.
+ """
+ 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 ''
+ 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]
+
+
+
+# โโ 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, 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)
+ # 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,
+ num_samples=1, logprobs=1,
+ temperature=args.keyword_temp, top_p=0.98))
+
+ from twinkle_agentic.protocol.openai import OpenAI
+ # 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,
+ 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.
+ self.system = P.SYSTEM + (P.BUILD_SIZE_CAP.format(n=args.max_build_files)
+ if args.max_build_files > 0 else '')
+ # 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
+ # 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()
+ # 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.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 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 = 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
+ # 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': 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.tools()}
+ 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
+
+ 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 = self.api.reply(messages, user_text, self.check_params)
+ 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 slot.run_script(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
+
+ 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
+ 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.tools()
+ attempt = rollout_one(self.solve_rollout, opening, self.solve_params, slot)
+ if attempt is not None:
+ 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
+ # 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.
+
+ 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 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):
+ self.keywords.remember_hard(group.keywords)
+ 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 {
+ '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())},
+ },
+ }
diff --git a/cookbook/rsi/agentic/episode.py b/cookbook/rsi/agentic/episode.py
new file mode 100644
index 000000000..07a19a96c
--- /dev/null
+++ b/cookbook/rsi/agentic/episode.py
@@ -0,0 +1,233 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Episode construction and scoring, shared by collection and eval.
+
+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).
+
+``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
+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.harness import MsAgentHarness
+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 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')
+ if setup:
+ 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:]}')
+
+ 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.tools()
+ trajectory['tools'] = schemas
+ tool_manager = env.tool_manager(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.run_script(check_script)
+ 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,
+ env=env,
+ )
+ 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..69d7bbda6
--- /dev/null
+++ b/cookbook/rsi/agentic/eval.py
@@ -0,0 +1,255 @@
+# 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 ``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 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 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
+ 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 \\
+ --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
+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',
+ 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('--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
+ # 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')])
+ 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('Template', model_id=args.model_id, enable_thinking=True,
+ max_length=args.max_model_len)
+ return sampler
+
+
+def main():
+ args = parse_args()
+ tasks = load_tasks(args.tasks)
+ cfg = SandboxConfig.from_env()
+ 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)
+ 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,
+ 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
+ # 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
new file mode 100644
index 000000000..9da9895c5
--- /dev/null
+++ b/cookbook/rsi/agentic/prompts.py
@@ -0,0 +1,427 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Every string this pipeline sends to a model.
+
+One proposal is one conversation with three stages:
+
+ 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``).
+
+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.
+
+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.
+"""
+
+# โโ 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']
+
+# 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 '
+ '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, 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. 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 '
+ '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
+ + _PINNED_EDGE,
+}
+
+# โโ 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.'
+)
+
+# 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, 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 '
+ '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.'
+)
+
+
+# โโ 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
+# 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 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: 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'
+ '- 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.'
+)
+
+# 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.'
+)
+
+# 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.\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 โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# 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.
+
+# 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'
+ "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'
+ '- 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'
+ '- 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.'
+)
+
+# โโ 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'
+ '{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'
+ '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 '
+ '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.'
+)
+
+# โโ 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.
+#
+# 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'
+ '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 you wrote.\n'
+ '- Everything DERIVED from it -- computed values, aggregates, orderings, '
+ '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'
+ '- 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.'
+)
diff --git a/cookbook/rsi/agentic/remote_tool_env.py b/cookbook/rsi/agentic/remote_tool_env.py
new file mode 100644
index 000000000..2e7b502cd
--- /dev/null
+++ b/cookbook/rsi/agentic/remote_tool_env.py
@@ -0,0 +1,630 @@
+"""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 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
+from twinkle.data_format.message import Tool as ToolInfo
+from twinkle_agentic.envs.base import Env, StepResult, truncate_observation
+
+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+)')
+
+# 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, io, traceback
+_tw_check_src = {body}
+_tw_buf = io.StringIO()
+_tw_out = sys.stdout
+sys.stdout = _tw_buf
+_tw_rc = 0
+try:
+ _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.
+ _c = _e.code
+ _tw_rc = 0 if _c is None else _c if isinstance(_c, int) else 1
+except BaseException:
+ 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')
+
+# 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:
+ """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.
+
+ 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: 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
+ 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
+ # 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
+
+ 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
+ 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'rebuilding the sandbox (recovery #{self.n_recoveries})')
+ self.reset()
+
+ 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.
+ """
+ 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 []
+ payload = {
+ 'calls': [{
+ 'tool_name': self._dispatch_name(name),
+ 'arguments': args or {}
+ } for name, args in calls],
+ 'timeout': timeout,
+ }
+ try:
+ 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
+ # 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
+ 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 tools(self) -> List[ToolInfo]:
+ """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._load_schemas()
+ return list(self._schemas)
+
+ def tool_names(self) -> List[str]:
+ names = []
+ for schema in self.tools():
+ name = (schema.get('function') or {}).get('name')
+ if name:
+ names.append(str(name))
+ return names
+
+ def resolve_tool(self, name: str) -> str:
+ """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.
+
+ 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
+ # 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]
+ 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 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.
+ """
+ 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.
+
+ 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:`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.
+ """
+ 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')
+ # ``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.
+
+ 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())
+
+ def _start_server(self) -> None:
+ """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'> {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:
+ 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.
+
+ 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.server_log(40)
+ 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.
+
+ 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 + _RPC_HEADROOM} http://127.0.0.1:{self._port}{path}'
+ else:
+ 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 + _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", "")}')
+ return json.loads(stdout)
+
+ def _truncate(self, text: str) -> str:
+ limit = self.max_observation_chars
+ return truncate_observation(text, limit) if limit else text
diff --git a/cookbook/rsi/agentic/rsi.py b/cookbook/rsi/agentic/rsi.py
new file mode 100644
index 000000000..224cea407
--- /dev/null
+++ b/cookbook/rsi/agentic/rsi.py
@@ -0,0 +1,305 @@
+# 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__)))
+_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()
+
+
+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_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, the sandbox pool and the recorder are owned by the caller and
+ outlive this.
+ """
+ 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.keywords.expand_hard()
+ finally:
+ 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. 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',
+ 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 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)
+ 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), 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 '
+ 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}')
+
+ # 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.
+ 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)
+ # 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
+ 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 ""}')
+
+ 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')
+ 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/rsi_agent.yaml b/cookbook/rsi/agentic/rsi_agent.yaml
new file mode 100644
index 000000000..430b7e920
--- /dev/null
+++ b/cookbook/rsi/agentic/rsi_agent.yaml
@@ -0,0 +1,190 @@
+# 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.
+ 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
+ # 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/run.sh b/cookbook/rsi/agentic/run.sh
new file mode 100644
index 000000000..d3b40740d
--- /dev/null
+++ b/cookbook/rsi/agentic/run.sh
@@ -0,0 +1,119 @@
+#!/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 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 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
+
+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}"
+# 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}"
+# 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,"
+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.
+#
+# 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" \
+ ${LIMIT[@]+"${LIMIT[@]}"} \
+ "$@" 2>&1 | tee -a "$LOG"
diff --git a/cookbook/rsi/agentic/sandbox.py b/cookbook/rsi/agentic/sandbox.py
new file mode 100644
index 000000000..99c309ec2
--- /dev/null
+++ b/cookbook/rsi/agentic/sandbox.py
@@ -0,0 +1,304 @@
+# 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 time
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
+from twinkle import get_logger
+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
+
+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(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: Env, 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._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)
+
+ # ------------------------------------------------------------------ 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.
+
+ 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():
+ 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_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.rebuild()
+ code, out = self.run_script(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_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_script(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 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].tools()
+ 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(s.env.n_recoveries 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/sandbox_server/Dockerfile b/cookbook/rsi/agentic/sandbox_server/Dockerfile
new file mode 100644
index 000000000..3a40cd35c
--- /dev/null
+++ b/cookbook/rsi/agentic/sandbox_server/Dockerfile
@@ -0,0 +1,103 @@
+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
+
+# 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.
+#
+# 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
+# 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.
+#
+# 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 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
new file mode 100644
index 000000000..d99300a49
--- /dev/null
+++ b/cookbook/rsi/agentic/sandbox_server/install.sh
@@ -0,0 +1,88 @@
+#!/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
+# 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/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
new file mode 100644
index 000000000..e24b13e02
--- /dev/null
+++ b/cookbook/rsi/agentic/sandbox_server/serve.sh
@@ -0,0 +1,120 @@
+#!/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
+# 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}"
+# 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}"
+
+# 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
+ 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 RUST_LOG=$RUST_LOG"
+
+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/rsi/agentic/sandbox_server/tool_server.py b/cookbook/rsi/agentic/sandbox_server/tool_server.py
new file mode 100644
index 000000000..2910abff6
--- /dev/null
+++ b/cookbook/rsi/agentic/sandbox_server/tool_server.py
@@ -0,0 +1,780 @@
+"""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 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, Set, Tuple
+
+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', )}
+
+_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 = '---'
+
+# 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.
+
+ 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.
+
+ 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.
+ """
+ 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):
+ 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)
+
+ setattr(python_executor, _SINGLE_NS_FLAG, True)
+ LocalCodeExecutionTool.python_executor = python_executor
+ 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.
+
+ 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:
+ 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 _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.
+
+ 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
+
+
+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.
+
+ 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
+
+ # 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
+ # 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.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.
+ self.agent.config.pop('llm', None)
+ self.agent._interactive = False
+ 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
+ # exists.
+ self._contracts = self._build_contracts()
+
+ 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
+
+ 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 [])
+ 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. 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*.
+ """
+ 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:
+ 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:
+ 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)
+
+
+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'
+
+ 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') or {}).get('name') for t in runtime.tools()]
+ llm_note = 'llm configured' if runtime.has_llm else 'no llm (read_file.abbreviate withdrawn)'
+ 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()
+
+
+if __name__ == '__main__':
+ main()
diff --git a/cookbook/rsi/agentic/tasks.example.jsonl b/cookbook/rsi/agentic/tasks.example.jsonl
new file mode 100644
index 000000000..d9d9a37e3
--- /dev/null
+++ b/cookbook/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/agentic/train.py b/cookbook/rsi/agentic/train.py
new file mode 100644
index 000000000..754baefc4
--- /dev/null
+++ b/cookbook/rsi/agentic/train.py
@@ -0,0 +1,355 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""One GRPO step on what a collection pass left in a run directory.
+
+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 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 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
+import os
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+
+from twinkle import DeviceMesh, get_logger
+from twinkle.advantage import GRPOAdvantage
+from twinkle.processor import InputProcessor
+
+logger = get_logger()
+
+
+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.
+ """
+ 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 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 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 = 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:
+ 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_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_length={max_length}'] += 1
+ continue
+ 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,
+ '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 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)}')
+ 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)
+ # 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)}; trainable tokens {dict(side_trainable)}')
+ 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))
+ # 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),
+ '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),
+ '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:
+ json.dump(summary, f, indent=2, ensure_ascii=False, default=str)
+ return summary
+
+
+# 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 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 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
+ 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/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'],
+ '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 step {iteration}: {len(log)} metrics')
diff --git a/cookbook/rsi/code/challenge.py b/cookbook/rsi/code/challenge.py
new file mode 100644
index 000000000..4e3b6383c
--- /dev/null
+++ b/cookbook/rsi/code/challenge.py
@@ -0,0 +1,274 @@
+# 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.envs import LocalEnv
+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__)))
+# 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()
+
+
+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-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)
+ 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 = MultiTurnRollout(
+ 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,
+ # 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,
+ 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_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,
+ )
+
+ 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')
+ # 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': check_script.splitlines(),
+ 'test_setup_code': ''}, ensure_ascii=False) + '\n')
+
+
+if __name__ == '__main__':
+ main()
diff --git a/cookbook/rsi/code/challenge_prompts.py b/cookbook/rsi/code/challenge_prompts.py
new file mode 100644
index 000000000..6e1388c49
--- /dev/null
+++ b/cookbook/rsi/code/challenge_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/code/collect.py b/cookbook/rsi/code/collect.py
new file mode 100644
index 000000000..2b15c6786
--- /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 MultiTurnRollout
+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 = 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)
+ 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/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/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
new file mode 100644
index 000000000..1e4062f0a
--- /dev/null
+++ b/cookbook/rsi/rl.py
@@ -0,0 +1,1208 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""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
+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.
+
+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.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 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)
+ 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
+ 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
+ (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 time
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Dict, List, Optional, Tuple
+
+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, 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()
+
+# โโ 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)
+# 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 (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))))
+# 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 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
+# 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 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
+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
+# 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', 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:
+ 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
+
+
+# โโ 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."""
+
+ 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
+
+
+# โโ 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. 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]]:
+ """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, unwrap_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'],
+ 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
+
+ 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', {})
+ 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 _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 the
+ flow 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 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).
+ """
+ tests = tests or {}
+ raw_by_query = raw_by_query or {}
+ 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 []
+ 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 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)
+ 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()
+ 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=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 _format_exec_error(err: str) -> str:
+ """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:]
+ 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 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.
+
+ 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.
+ """
+ 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
+ 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.
+ """
+ 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,
+ 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, MODEL_GPUS + SAMPLER_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)
+ 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)
+
+ # 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.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={
+ '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, 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 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)
+
+ 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')
+
+ 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)')
+ 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)
+
+ # No LoRA, so every sync ships the full weights.
+ ckpt_manager.sync_weights(merge_and_sync=True)
+ sampler.reset_prefix_cache()
+
+ all_tokens: List[List[int]] = []
+ if SOLVER_MODE == 'grpo':
+ # 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 = [], [], []
+ 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)
+ 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)
+ 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'{SAVE_NAME}-checkpoint-{optim_step}', output_dir=SAVE_DIR)
+
+ log_dict = metrics.calculate()
+ 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(SAVE_NAME, output_dir=SAVE_DIR)
+
+
+if __name__ == '__main__':
+ main()
diff --git a/cookbook/rsi/run_rsi.py b/cookbook/rsi/run_rsi.py
new file mode 100644
index 000000000..f1b9ebeca
--- /dev/null
+++ b/cookbook/rsi/run_rsi.py
@@ -0,0 +1,148 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""RSI entry point โ run one stage on its own, or the whole chain.
+
+Three stages, one model (Qwen3-4B) playing both roles:
+
+ 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:
+ * ``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 wires each stage's default output into
+the next stage's input and forwards any extra flags straight through.
+
+Examples
+--------
+ # clean a dataset into seeds (optional)
+ python cookbook/rsi/run_rsi.py --step prepare --raw data/raw.jsonl
+
+ # 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
+import os
+import subprocess
+import sys
+
+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', 'challenge', 'rl']
+
+
+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 one.
+ """
+ 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 (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.seeds, '--num-proc', str(a.num_proc)]
+ if a.dropped_log:
+ argv += ['--dropped-log', a.dropped_log]
+ return argv + 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':
+ # 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] ๆช็ฅ step: {step}')
+
+
+def main():
+ parser = argparse.ArgumentParser(
+ 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->challenge->rl).')
+ parser.add_argument('--raw', default='', help='Raw data source for prepare (local path or ms:// id).')
+ 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:
+ argv, env = _argv_for(step, a, [])
+ _run(SCRIPTS[step], argv, env)
+ print('\n[run_rsi] all stages done.', flush=True)
+ return
+
+ argv, env = _argv_for(a.step, a, extra)
+ _run(SCRIPTS[a.step], argv, env)
+
+
+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: <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 \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: <trigger>`; numbered telegraphic \
+steps `1.do X 2.then Y`; close with `Output: <result>` 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: <subject \u2014 scope>`, 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 `- <topic-key>: <one-line hint>`:
+ \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: <specific pattern name \u2014 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.')
+
+# ---------------------------------------------------------------------------
+# 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/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 94b143454..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,9 +76,10 @@ 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 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
@@ -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 b74c1e791..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,9 +76,10 @@ 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` | ่ฝจ่ฟนๆฏๅฆ่ขซๆชๆญใ |
+| `truncated` | `bool` | ่ฝจ่ฟนๆฏๅฆ่ขซๆชๆญ๏ผ่้่ช่ก็ปๆ๏ผ๏ผ็ๆ่งฆๅ `max_tokens`๏ผ`stop_reason='length'`๏ผใ่พพๅฐ่ฝฎๆฌกไธ้๏ผๆ่ขซ้ฟๅบฆไธ้ไธขๅผใ |
| `logprobs` | `List` | ๆฏ token ็ๅฏนๆฐๆฆ็๏ผๅฆๆ๏ผใ |
### Ray ่ฟ็จๆฏๆ
@@ -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/pyproject.toml b/pyproject.toml
index f1b8df372..b9d799460 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/data_format/__init__.py b/src/twinkle/data_format/__init__.py
index 1dff273c7..5db25a2b5 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, SamplingMask, SamplingParams
-from .trajectory import Trajectory, 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/sampling.py b/src/twinkle/data_format/sampling.py
index cdd2233a8..a1095dfe3 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/data_format/trajectory.py b/src/twinkle/data_format/trajectory.py
index 992df28d1..c7b277a35 100644
--- a/src/twinkle/data_format/trajectory.py
+++ b/src/twinkle/data_format/trajectory.py
@@ -1,7 +1,8 @@
# Copyright (c) ModelScope Contributors. All rights reserved.
import json
import sys
-from typing import Any, List, Optional, Tuple, Union
+from collections.abc import Mapping
+from typing import Any, Dict, List, Optional, Tuple, Union
from .message import Message, Tool
@@ -28,8 +29,44 @@ 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 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):
+ 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/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 281e020dd..ffd4e1a42 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 = {
@@ -367,6 +373,9 @@ def create_workers(worker_cls: Type[T],
# (for example, 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:
@@ -390,6 +399,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 = {
@@ -401,6 +415,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/loss/__init__.py b/src/twinkle/loss/__init__.py
index 18b93f5ae..926dcb8a3 100644
--- a/src/twinkle/loss/__init__.py
+++ b/src/twinkle/loss/__init__.py
@@ -8,6 +8,7 @@
from .infonce import InfonceLoss
from .liger_fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
from .mse import MSELoss
+from .opsd import OPSDLoss
from .value import PPOValueLoss
torch_loss_mapping = {
@@ -26,6 +27,8 @@
'cispo': CISPOLoss,
'bnpo': BNPOLoss,
'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 56f56ebec..85470a0dc 100644
--- a/src/twinkle/loss/grpo.py
+++ b/src/twinkle/loss/grpo.py
@@ -73,8 +73,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(
@@ -84,9 +84,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
@@ -100,7 +107,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,
@@ -126,6 +140,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]',
@@ -202,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,
@@ -244,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]:
@@ -315,7 +373,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 PPOLoss(GRPOLoss):
@@ -367,7 +425,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
@@ -456,20 +514,58 @@ class BNPOLoss(GRPOLoss):
BNPO (Batch-Normalized Policy Optimization) Loss.
Normalizes by total completion tokens across batch.
+
+ ``token_mean_scope``:
+ '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 <think> 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 = '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 micro_batch_scale(self, inputs, indices):
- token_counts = []
- for model_input in inputs:
- labels = model_input['labels']
- if hasattr(labels, 'ne'):
- token_counts.append(int(labels.ne(self.ignore_index).sum().item()))
- else:
- token_counts.append(sum(int(token != self.ignore_index) for token in labels))
- total_tokens = sum(token_counts)
- if total_tokens == 0:
- return 0.0
- return sum(token_counts[index] for index in indices) / total_tokens
+ """The weight one micro-batch carries, which has to follow ``token_mean_scope``.
+
+ 'global' already returns the token SUM and reports ``num_tokens=ฮฃmask``, so the
+ global division happens downstream; scaling here as well would divide twice. Same
+ contract as ``CrossEntropyLoss(reduction='sum')``.
+
+ 'micro' *is* the equal-weighted mean of per-micro token-means, so the inherited
+ sample fraction is its weight. A token fraction would make the micro losses sum to
+ the global token-mean and erase the distinction this scope exists to make -- which
+ is the +3.2e-4 vs +0.031 per-token pg_loss measured above.
+ """
+ if self.token_mean_scope == 'global':
+ return 1.0
+ return super().micro_batch_scale(inputs, indices)
def _aggregate_loss(
self,
@@ -477,8 +573,18 @@ def _aggregate_loss(
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 DRGRPOLoss(GRPOLoss):
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/metric/grpo.py b/src/twinkle/metric/grpo.py
index e3eaacd20..12c978f3e 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))
@@ -288,25 +339,38 @@ 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
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 +390,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 +411,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 +440,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/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/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())
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 cf9c7c405..9ce4c4cb7 100644
--- a/src/twinkle/model/transformers/transformers.py
+++ b/src/twinkle/model/transformers/transformers.py
@@ -262,6 +262,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.
@@ -563,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(
@@ -590,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)
@@ -684,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'))
@@ -715,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/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/src/twinkle/preprocessor/__init__.py b/src/twinkle/preprocessor/__init__.py
index 58e3a2b05..c29e0f296 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,
DAPOMathProcessor, GSM8KProcessor, SelfCognitionProcessor)
diff --git a/src/twinkle/preprocessor/base.py b/src/twinkle/preprocessor/base.py
index 0225d3c1e..4d9cb51c0 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]]:
@@ -29,15 +42,58 @@ 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
- 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 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:
+ if self.keep(r):
+ kept.append(r)
+ else:
+ dropped.append(dict(r, drop_reason=self.drop_reason(r)))
+ return kept, dropped
class DataFilter:
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 f4a27cb97..7877ddb55 100644
--- a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py
+++ b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py
@@ -36,7 +36,10 @@ def _convert_ndarray_to_list(obj: Any) -> Any:
return obj
-@remote_class()
+_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).
@@ -278,7 +281,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]],
@@ -493,49 +496,77 @@ def unload_adapter_paths(self, adapter_paths: list[str]) -> None:
self._run_in_loop(self.engine.unload_lora_paths(adapter_paths))
@remote_function(dispatch='all', collect='first', lazy_collect=False)
- def load_full_weights_from_path(self, path: str) -> int:
+ def load_full_weights_from_path(self, path: Optional[str] = None) -> int:
"""Load a full (non-LoRA) HF checkpoint into the engine's base model.
- Used by full-parameter training: the saved checkpoint is a plain HF
- directory (no ``adapter_config.json``), so it replaces the sampler's
- base weights instead of being loaded as a LoRA adapter. Idempotent:
- repeated calls with the same resolved path are skipped.
+ Unlike :meth:`receive_weights`, this does **not** involve the training model:
+ weights are read from disk and streamed straight into vLLM. That is what lets
+ a sampler be restored 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. Full-parameter training uses the same entry point:
+ its checkpoint is a plain HF directory (no ``adapter_config.json``), so it
+ replaces the base weights instead of loading as a LoRA adapter.
+
+ 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).
+
+ Idempotent: repeated calls with the same resolved path are skipped.
+
+ 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.
Returns:
- 1 if weights were (re)loaded, 0 if the path was already loaded.
+ 1 if weights were (re)loaded, 0 if that path was already loaded.
"""
import glob
import json
- import os
+ from safetensors import safe_open
- resolved = HubOperation.download_model(model_id_or_path=path)
+ path = path or self.model_id
+ resolved = path if os.path.exists(path) else HubOperation.download_model(path)
if getattr(self, '_loaded_full_weights_path', None) == resolved:
return 0
- from safetensors import safe_open
-
- def _weight_iter():
- index = os.path.join(resolved, 'model.safetensors.index.json')
- if os.path.exists(index):
- with open(index) as f:
- shards = sorted(set(json.load(f)['weight_map'].values()))
- files = [os.path.join(resolved, s) for s in shards]
- else:
- files = sorted(glob.glob(os.path.join(resolved, '*.safetensors')))
- for fp in files:
- with safe_open(fp, framework='pt', device='cpu') as f:
- for key in f.keys():
- yield key, f.get_tensor(key)
+ # Resolve the shard list eagerly so a bad path fails here rather than
+ # part-way through streaming tensors into a live engine.
+ index_path = os.path.join(resolved, '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(resolved, s) for s in sorted(set(weight_map.values()))]
+ else:
+ shards = sorted(glob.glob(os.path.join(resolved, '*.safetensors')))
+ if not shards:
+ raise FileNotFoundError(f'No .safetensors weights found under {resolved}')
+
+ 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)
async def _load():
- await self.engine.update_weights(_weight_iter(), peft_config=None, base_sync_done=False)
- # A full base-model swap invalidates any previously synced LoRA.
+ await self.engine.update_weights(_iter_weights(), peft_config=None, base_sync_done=False)
+ # A base-model swap invalidates any previously synced LoRA adapter,
+ # mirroring the `not base_sync_done` branch of receive_weights().
self.engine.invalidate_synced_lora()
logger.info(f'Loading full-parameter weights into sampler base model from {resolved}')
self._run_in_loop(_load())
self._loaded_full_weights_path = resolved
+ # Prefixes cached under the previous weights would decode against a model
+ # that no longer exists; drop them before the next sample().
self.reset_prefix_cache()
+ logger.info(f'Reloaded base weights from {resolved} ({len(shards)} shard(s))')
return 1
@remote_function(dispatch='all', collect='first', lazy_collect=False)
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 3c6c29f6c..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:
@@ -87,6 +99,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
@@ -187,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):
@@ -217,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
@@ -226,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':
@@ -264,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):
@@ -311,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))
@@ -340,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:
@@ -514,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
@@ -526,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 [])
@@ -682,7 +772,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]):
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/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
new file mode 100644
index 000000000..a9facc622
--- /dev/null
+++ b/src/twinkle/template/tools/bracket_dsl.py
@@ -0,0 +1,256 @@
+# 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.
+
+ 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'
+ 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)
+ # 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:
+ # 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]]:
+ """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]] = []
+ skip = BracketDslParser._skip_spans(text)
+ i, n = 0, len(text or '')
+ while i < n:
+ if text[i] != '[' or BracketDslParser._in_spans(i, skip):
+ 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:
+ 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)
+ 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]]:
+ 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]
+ 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:
+ errors.append(f'{m.group(1).strip()}( is never closed by a '
+ f'matching )')
+ break
+ name = m.group(1).strip()
+ if name:
+ calls.append({
+ 'type': 'function',
+ 'function': {
+ 'name': name,
+ '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, errors
+
+ 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/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/utils/__init__.py b/src/twinkle/utils/__init__.py
index 164cd1561..e983ef29f 100644
--- a/src/twinkle/utils/__init__.py
+++ b/src/twinkle/utils/__init__.py
@@ -15,5 +15,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/torch_utils.py b/src/twinkle/utils/torch_utils.py
index 42bdaa0c9..88a787ecc 100644
--- a/src/twinkle/utils/torch_utils.py
+++ b/src/twinkle/utils/torch_utils.py
@@ -408,17 +408,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/src/twinkle/utils/transformers_utils.py b/src/twinkle/utils/transformers_utils.py
index 9f997bf70..674e019bb 100644
--- a/src/twinkle/utils/transformers_utils.py
+++ b/src/twinkle/utils/transformers_utils.py
@@ -22,6 +22,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
@@ -48,8 +62,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/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/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/__init__.py b/src/twinkle_agentic/challenger/__init__.py
new file mode 100644
index 000000000..0719b04e5
--- /dev/null
+++ b/src/twinkle_agentic/challenger/__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/agentic.py b/src/twinkle_agentic/challenger/agentic.py
new file mode 100644
index 000000000..ca41486a7
--- /dev/null
+++ b/src/twinkle_agentic/challenger/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/base.py b/src/twinkle_agentic/challenger/base.py
new file mode 100644
index 000000000..1c819533b
--- /dev/null
+++ b/src/twinkle_agentic/challenger/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/keyword.py b/src/twinkle_agentic/challenger/keyword.py
new file mode 100644
index 000000000..09dee2bad
--- /dev/null
+++ b/src/twinkle_agentic/challenger/keyword.py
@@ -0,0 +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 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 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 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/recorder.py b/src/twinkle_agentic/challenger/recorder.py
new file mode 100644
index 000000000..c72d6fab3
--- /dev/null
+++ b/src/twinkle_agentic/challenger/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/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
- ``<block_N>``.
- """
-
- 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/condenser/__init__.py b/src/twinkle_agentic/condenser/__init__.py
deleted file mode 100644
index e78545002..000000000
--- a/src/twinkle_agentic/condenser/__init__.py
+++ /dev/null
@@ -1,5 +0,0 @@
-from .base import Condenser
-from .keyword import KeywordCondenser
-from .model import ModelCondenser
-
-__all__ = ['Condenser', 'KeywordCondenser', 'ModelCondenser']
diff --git a/src/twinkle_agentic/condenser/base.py b/src/twinkle_agentic/condenser/base.py
deleted file mode 100644
index 5e42dab17..000000000
--- a/src/twinkle_agentic/condenser/base.py
+++ /dev/null
@@ -1,13 +0,0 @@
-from abc import ABC, abstractmethod
-
-from twinkle_agentic.data_format import Chunks
-
-
-class Condenser(ABC):
- """
- TODO: Experimental feature, wait for testing
- """
-
- @abstractmethod
- def __call__(self, chunks: Chunks, **kwargs) -> Chunks:
- raise NotImplementedError
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 ``<block_N>...</block_N>``.
-
- 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|>``, ``</s>``,
- ...) 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 ``<block_N>``.
-
- 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 ``<block_N>`` 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 (``</s>``) 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/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 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 552d9e1a5..52df86ce9 100644
--- a/src/twinkle_agentic/envs/base.py
+++ b/src/twinkle_agentic/envs/base.py
@@ -1,11 +1,104 @@
# 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
+# 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:
@@ -24,8 +117,19 @@ 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.
"""
+ #: 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()
@@ -33,9 +137,123 @@ 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 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/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/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/__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..dadabce63
--- /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 = 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.
+ """
+ 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..c9da5774c
--- /dev/null
+++ b/src/twinkle_agentic/harness/ms_agent.py
@@ -0,0 +1,557 @@
+# 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::
+
+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
+**not** call ``llm.generate`` or ``parallel_tool_call`` (those execute tools).
+"""
+from __future__ import annotations
+
+import json
+import os
+import sys
+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
+ # 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
+ 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:
+ # 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()
+ 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
+
+ # _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)
+
+ 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 ''
+ 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
+
+ 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."""
+ patch_ms_agent_python_executor()
+ 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()
+ 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()
+ 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
+
+
+_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.
+
+ 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.
+ """
+ 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):
+ 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)
+
+ 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 []
+ 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/preprocessor/__init__.py b/src/twinkle_agentic/preprocessor/__init__.py
index a69b44392..9f5deb24c 100644
--- a/src/twinkle_agentic/preprocessor/__init__.py
+++ b/src/twinkle_agentic/preprocessor/__init__.py
@@ -12,18 +12,80 @@
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 .pii_presidio_filter import PIIPresidioFilter
from .refuse_filter import RefuseFilter
-from .score_filter import ScoreFilter
from .token_soup import TokenSoupFilter
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.<pid>`` 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.
@@ -31,19 +93,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:
@@ -59,12 +147,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/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 75cf3d00b..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 โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
#
@@ -155,6 +156,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 <think> reasoning (free to ramble),
@@ -175,6 +177,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)
@@ -189,18 +192,29 @@ 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
- 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/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/hard_filter.py b/src/twinkle_agentic/preprocessor/hard_filter.py
index 043389ca0..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, msg_content_text, msg_has_media
+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 โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
@@ -82,14 +83,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 +124,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 +138,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 +197,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..6d1b21c24 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()
@@ -205,10 +205,25 @@ class _RegexDetector(IntentDetector):
"""Common scaffolding: scan messages, run ``_match`` on each text, pair to assistant."""
role_filter: Optional[str] = None
+ # Whether ``<think>`` 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
+ # <think> 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
+ # <think>...</think> 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):
@@ -221,7 +236,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)
@@ -245,6 +260,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
@@ -258,6 +274,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
@@ -269,6 +286,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
@@ -337,6 +355,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 +366,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/language_filter.py b/src/twinkle_agentic/preprocessor/language_filter.py
new file mode 100644
index 000000000..b86c0b021
--- /dev/null
+++ b/src/twinkle_agentic/preprocessor/language_filter.py
@@ -0,0 +1,114 @@
+# 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 twinkle_agentic.utils.message_utils import msg_content_text
+from twinkle_agentic.utils.text_utils import cjk_ratio
+
+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[^>]*>.*?</\1>',
+ re.DOTALL | re.IGNORECASE,
+)
+# Self-closing / unmatched openers of the same tags (defensive).
+_INJECTION_TAG_RE = re.compile(
+ r'</?(system-reminder|system_reminder|system|instructions?|context|'
+ r'important_instructions|env|environment|tools?)\b[^>]*/?>',
+ 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
+
+ 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/llm_backend.py b/src/twinkle_agentic/preprocessor/llm_backend.py
deleted file mode 100644
index 002618620..000000000
--- a/src/twinkle_agentic/preprocessor/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/message_normalizer.py b/src/twinkle_agentic/preprocessor/message_normalizer.py
index d3074a565..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
+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".
@@ -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:
@@ -197,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:
@@ -206,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/message_sanity.py b/src/twinkle_agentic/preprocessor/message_sanity.py
index 6001e1d7e..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_.\-]*$')
@@ -31,6 +28,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 +36,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 +338,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/model_filter.py b/src/twinkle_agentic/preprocessor/model_filter.py
index fe238b1ed..4162c54c5 100644
--- a/src/twinkle_agentic/preprocessor/model_filter.py
+++ b/src/twinkle_agentic/preprocessor/model_filter.py
@@ -1,10 +1,21 @@
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.
+#
+# 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',
@@ -21,7 +32,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 +40,8 @@ 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 ''))
+
+ def drop_reason(self, row: Dict[str, Any]) -> str:
+ return 'model_not_allowed'
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/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'<think(?:ing)?>.*?</think(?:ing)?>\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/score_filter.py b/src/twinkle_agentic/preprocessor/score_filter.py
deleted file mode 100644
index 8530ab119..000000000
--- a/src/twinkle_agentic/preprocessor/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 .llm_backend import LLMBackend
-from .utils import _chr_min_distinct, _ifd_family_metrics, _lp_to_jsonable, _pad_batch, _to_int_list
-
-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/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 7a041fcec..000000000
--- a/src/twinkle_agentic/preprocessor/utils.py
+++ /dev/null
@@ -1,359 +0,0 @@
-"""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
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/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/rollout/__init__.py b/src/twinkle_agentic/rollout/__init__.py
index 67c589e06..cddff8eb2 100644
--- a/src/twinkle_agentic/rollout/__init__.py
+++ b/src/twinkle_agentic/rollout/__init__.py
@@ -1,14 +1,7 @@
# Copyright (c) ModelScope Contributors. All rights reserved.
-from .api_multi_turn import APIMultiTurnRollout
+from .api_sampler import APISampler
from .base import Rollout
from .bridge import extend_with_bridge
from .multi_turn import MultiTurnRollout
-from .multi_turn_condense import MultiTurnCondenseRollout
-__all__ = [
- 'APIMultiTurnRollout',
- 'MultiTurnCondenseRollout',
- 'MultiTurnRollout',
- 'Rollout',
- 'extend_with_bridge',
-]
+__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 7521a454f..000000000
--- a/src/twinkle_agentic/rollout/api_multi_turn.py
+++ /dev/null
@@ -1,289 +0,0 @@
-# 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.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'
-_STOP_LENGTH = 'length'
-_STOP_MAX_TURNS = 'max_turns'
-_STOP_API_ERROR = 'api_error'
-
-
-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``).
-
- 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.
-
- 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: OpenAI,
- tool_manager: ToolManager,
- 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 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}')
- 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.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,
- 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 = MultiTurnRollout._resolve_tool_managers(kwargs.get('tool_manager', self.tool_manager), n)
- extra_body = dict(self.extra_body)
- if 'extra_body' in kwargs and kwargs['extra_body']:
- extra_body.update(kwargs['extra_body'])
-
- # 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): 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_traces(result_outs, kwargs.get('global_step'))
- return result_outs
-
- # ------------------------------------------------------------------ private
-
- def _run_one(
- self,
- trajectory: Trajectory,
- tool_manager: ToolManager,
- sampling_params: SamplingParams,
- extra_body: Dict[str, Any],
- ) -> 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:
- 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
-
- 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
-
- out = dict(trajectory)
- out['messages'] = messages
- out['turns'] = turn
- out['stop_reason'] = stop_reason
- out['truncated'] = truncated
- 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 _write_traces(
- 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
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 64d9f922b..dfaa9ae50 100644
--- a/src/twinkle_agentic/rollout/base.py
+++ b/src/twinkle_agentic/rollout/base.py
@@ -1,12 +1,295 @@
# 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 concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Any, Callable, Dict, List, Optional, Tuple
-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
+
+# 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.
+
+ 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.
+
+ 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
+ # 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
+ concurrency: Optional[int] = None
+
+ # ------------------------------------------------------------------ setup
+
+ def _init_common(
+ self,
+ *,
+ 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,
+ ) -> 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}')
+ 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
- @abstractmethod
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 _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:
+ 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)
+ 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
+
+ _TRACE_SKIP_KEYS = (
+ 'input_ids',
+ 'labels',
+ 'completion_mask',
+ '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/bridge.py b/src/twinkle_agentic/rollout/bridge.py
index 2663f9ed1..ce668d54b 100644
--- a/src/twinkle_agentic/rollout/bridge.py
+++ b/src/twinkle_agentic/rollout/bridge.py
@@ -1,22 +1,35 @@
# 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.
"""
+
+
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.
@@ -40,6 +53,101 @@ def _to_plain(obj: Any) -> Any:
return obj
+def _delta_text(
+ template: Template,
+ 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)
+
+ 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)
+
+ 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
+ # 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 = 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 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
+ 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]],
@@ -63,31 +171,16 @@ def extend_with_bridge(
Returns ``None`` when the trajectory exceeds ``max_length`` and the
template's truncation strategy is ``'delete'``.
"""
- tokenizer = template.tokenizer
-
messages_before = list(pif.get('messages') or [])
messages_after = messages_before + list(tool_messages)
- 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)
-
- 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}')
- bridge_text = s_after[len(s_before):]
+ 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}')
@@ -107,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
@@ -129,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/multi_turn.py b/src/twinkle_agentic/rollout/multi_turn.py
index 786e4261b..6c9fd03d2 100644
--- a/src/twinkle_agentic/rollout/multi_turn.py
+++ b/src/twinkle_agentic/rollout/multi_turn.py
@@ -1,383 +1,764 @@
# Copyright (c) ModelScope Contributors. All rights reserved.
import json
-import os
import re
-import time
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Literal, Optional, Tuple
-from twinkle.data_format import Trajectory, user_data_get
-from twinkle.data_format.sampling import SampleResponse, SamplingParams
+from twinkle.data_format import Trajectory
+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]],
+) -> 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 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],
+) -> 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
+
+
+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).
+ """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:
- 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
+ response_callback(...) # sampler or API -> SampledSequence
+ harness.after_generate
+ 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.
- * ``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).
+ * ``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 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 max_turns < 1:
- raise ValueError(f'max_turns must be >= 1, got {max_turns}')
+ 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.sampling_params = sampling_params or SamplingParams()
- self.max_turns = max_turns
+ 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
+ # 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
- 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}')
+ # 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
+ # 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
+ # 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.')
@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)
+
+ 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 {}
+ 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)
- tool_managers = self._resolve_tool_managers(kwargs.get('tool_manager', self.tool_manager), n)
-
- # 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:
- pif = self.template.encode(traj, add_generation_prompt=True)
- pif = _to_plain(pif)
- pif.setdefault('messages', list(traj.get('messages', [])))
- pifs.append(pif)
-
- 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
-
- for _ in range(self.max_turns):
- active = [i for i in range(n) if not done[i]]
- if not active:
+ 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).
+ 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 = 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 pif, live, followups, budget, params, truncated
+ if followup_fn is None or followups >= MAX_FOLLOWUPS:
+ return False
+ followup = followup_fn(
+ self._as_trajectory(trajectory, pif, logprobs, turns, stop_reason, truncated), followups)
+ if followup is None:
+ return False
+ text, next_params = self._unpack_followup(followup)
+ extended = extend_with_bridge(pif, [{'role': 'user', 'content': text}], self.template)
+ if extended is None:
+ truncated = True
+ return False
+ 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 = next_params
+ return True
+
+ while spent < budget:
+ spent += 1
+
+ if spent > 1:
+ pif, live, dropped = self._harness_before_generate(pif, live, harness)
+ if dropped:
+ truncated = True
+ break
+
+ # 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
- # 2. One batched sample call for all currently-live trajectories.
- batch_pifs = [pifs[i] for i in active]
- actual = len(batch_pifs)
- 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]
-
- pending_bridges: List[tuple] = [] # (global_idx, tool_messages)
- for local_idx, global_idx in enumerate(active):
- turns[global_idx] += 1
- seq = resps[local_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.')
-
- 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
-
- # 3. Termination conditions
- if seq.stop_reason == 'length':
- done[global_idx] = True
- continue
+ # 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
- # 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
+ 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
+ pif = extended
+ if live is not None:
+ live['messages'] = list(extended.get('messages') or [])
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
+ # 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 turns[global_idx] >= self.max_turns:
- truncated[global_idx] = True
- 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
- # 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 = 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
+ 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:
- pifs[global_idx] = extended
-
- 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]
- 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
+ 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
+
+ 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 _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
-
- _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)
+ 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.
- @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 ''
+ 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 _resolve_traj_id(traj: Dict[str, Any], fallback_idx: int) -> str:
- """Stable-ish trajectory id for filenames.
+ 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.
- 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.
+ 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.
"""
- 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(
+ 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,
- 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,
- }
+ 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 = extend_with_bridge(pif, delta, self.template)
+ if extended is None:
+ return pif, live, True
+ live['messages'] = list(extended.get('messages') or [])
+ return extended, live, False
- 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.
+ @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.
"""
- if not self.trace_dir:
+ pif_msgs = pif.get('messages') or []
+ if not pif_msgs or pif_msgs[-1].get('role') != 'assistant':
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
+ 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]
- @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
+ 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
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/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/src/twinkle_agentic/summarizer/__init__.py b/src/twinkle_agentic/summarizer/__init__.py
new file mode 100644
index 000000000..53d48ff21
--- /dev/null
+++ b/src/twinkle_agentic/summarizer/__init__.py
@@ -0,0 +1,3 @@
+from .base import Summarizer
+
+__all__ = ['Summarizer']
diff --git a/src/twinkle_agentic/summarizer/base.py b/src/twinkle_agentic/summarizer/base.py
new file mode 100644
index 000000000..aecc25675
--- /dev/null
+++ b/src/twinkle_agentic/summarizer/base.py
@@ -0,0 +1,216 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from __future__ import annotations
+
+import math
+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
+
+
+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 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``.
+ 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:
+ 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.
+ 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: 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. 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,
+ backend: Any,
+ 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,
+ **rollout_kwargs: Any,
+ ):
+ 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:
+ 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.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
+ # 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)
+ # ------------------------------------------------------------------
+ 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.
+
+ 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
+ # ------------------------------------------------------------------
+ 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._backend, '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 = 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
+ return text
+
+ @staticmethod
+ def _strip_code_fences(text: str) -> str:
+ stripped = text.strip()
+ m = Summarizer._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/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..2bdf01410 100644
--- a/src/twinkle_agentic/tools/tool_manager.py
+++ b/src/twinkle_agentic/tools/tool_manager.py
@@ -1,11 +1,15 @@
# 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
+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."""
@@ -19,6 +23,59 @@ 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__}.')
+
+
+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__(
@@ -65,33 +122,112 @@ 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__}.')
-
+ 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
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 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.
+
+ 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:
+ return []
+ if len(calls) == 1:
+ return [self(calls[0])]
+
+ unpacked = [_unpack_tool_call(tc) for tc in calls]
+ env = self._shared_env()
+ 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 _i, name, args in batched])
+ 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)
+ 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/utils/code_utils.py b/src/twinkle_agentic/utils/code_utils.py
new file mode 100644
index 000000000..ab8f38403
--- /dev/null
+++ b/src/twinkle_agentic/utils/code_utils.py
@@ -0,0 +1,111 @@
+# 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: Optional[Tuple[str, ...]]) -> Pattern:
+ """Match a fenced block, optionally restricting its language label.
+
+ ``None`` accepts any label. Otherwise, listed tags match case-insensitively,
+ with any version suffix; an unlabelled fence is accepted as well.
+ """
+ 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)
+
+
+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: Optional[Tuple[str, ...]] = PYTHON_TAGS,
+) -> Optional[str]:
+ """Return the last matching fenced block, or None if there is none.
+
+ 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
+ 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/utils/llm_backup.py b/src/twinkle_agentic/utils/llm_backup.py
new file mode 100644
index 000000000..a7d804101
--- /dev/null
+++ b/src/twinkle_agentic/utils/llm_backup.py
@@ -0,0 +1,357 @@
+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
+ trajectory: Any = None
+
+
+@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,
+ 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:
+ 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()
+_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():
+ """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
+ # 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', '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
+
+
+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, trajectory=trajectory)
+ 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, trajectory=trajectory)
+ 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, 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, trajectory=trajectory)
+ 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
diff --git a/src/twinkle_agentic/utils/message_utils.py b/src/twinkle_agentic/utils/message_utils.py
new file mode 100644
index 000000000..64a543b6c
--- /dev/null
+++ b/src/twinkle_agentic/utils/message_utils.py
@@ -0,0 +1,127 @@
+# 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
+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:
+ """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))
+
+
+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
+
+
+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
+
+
+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
new file mode 100644
index 000000000..5dd8cc2ac
--- /dev/null
+++ b/src/twinkle_agentic/verifier/__init__.py
@@ -0,0 +1,12 @@
+from .result_check import (Check, CheckContext, CheckOutcome, CheckReport,
+ 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',
+ '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
new file mode 100644
index 000000000..d619a8002
--- /dev/null
+++ b/src/twinkle_agentic/verifier/result_check.py
@@ -0,0 +1,318 @@
+# 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``) 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
+from dataclasses import dataclass, field
+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',
+ 'CheckOutcome',
+ 'CheckReport',
+ 'CheckContext',
+ 'run_checks',
+ 'checks_from_dicts',
+]
+
+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_GB = 2.0
+
+_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.
+ env: where the ``shell`` / ``python`` kinds run -- the environment the
+ episode acted in. ``None`` runs them locally in ``workspace``.
+ """
+ workspace: str = ''
+ final_answer: str = ''
+ env: Optional['Env'] = None
+
+
+def _local_env(workspace: str) -> 'Env':
+ """Run checks in ``workspace`` on this machine.
+
+ 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.
+ """
+ # 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]:
+ """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'):
+ env = ctx.env or _local_env(ctx.workspace)
+ try:
+ code, out = env.run_script(check.code, kind, check.timeout)
+ except Exception as e: # noqa
+ 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:]}')
+ 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 / environment.
+ 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_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))
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/loss/test_bnpo_token_mean.py b/tests/loss/test_bnpo_token_mean.py
new file mode 100644
index 000000000..0ca009947
--- /dev/null
+++ b/tests/loss/test_bnpo_token_mean.py
@@ -0,0 +1,87 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""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
+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_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 # 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_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_matches_verl_equal_weighted_micro_means()
+ test_global_reports_sum_reduction_for_display()
+ print('OK: micro (default) == verl equal-weighted micro-means; global is split-invariant')
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'])
diff --git a/tests/metric/test_metrics.py b/tests/metric/test_metrics.py
index 691da2c06..322235d66 100644
--- a/tests/metric/test_metrics.py
+++ b/tests/metric/test_metrics.py
@@ -354,6 +354,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/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 <think>.
+
+Regression for the copywriting-tagged-as-code bug: a non-code answer whose
+private <think> 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('<think>1. ๅๆ้ๆฑ\n```\n้ฉๅญโ็็นโ่ฝฌๅ\n```\n</think>'
+ '้ฉๅญ๏ผ่ฟไปทๆ ผๆไธๆขไฟกใ่ฝฌๅ๏ผ็ฐๅจไธๅ็ซ็ๅ
ซๅใ'),
+ ]
+ assert CodeDetector()(msgs) == []
+
+
+def test_real_code_in_visible_answer_still_detected():
+ msgs = [
+ {'role': 'user', 'content': 'ๅไธชๅฟซๆ'},
+ _asst('<think>ๅ
ๆณ่พน็</think>```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'<think>ๅฏไปฅ็จ \frac{a}{b} \sum \int \sqrt{x} ๆไธชๆฏๆน</think>'
+ 'ๆถฆ่ฒๅ็ๆๅญ๏ผ้้กบ่ช็ถใ'),
+ ]
+ assert MathDetector()(msgs) == []
diff --git a/tests/preprocessor/test_preprocessor_utils.py b/tests/preprocessor/test_preprocessor_utils.py
deleted file mode 100644
index d52f8a77f..000000000
--- a/tests/preprocessor/test_preprocessor_utils.py
+++ /dev/null
@@ -1,354 +0,0 @@
-# Copyright (c) ModelScope Contributors. All rights reserved.
-"""Tests for preprocessor.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.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/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_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/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
new file mode 100644
index 000000000..37e96d584
--- /dev/null
+++ b/tests/twinkle_agentic/test_agentic_rsi.py
@@ -0,0 +1,1708 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""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
+import re
+import shutil
+import sys
+import tempfile
+import threading
+import unittest
+
+_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', '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)
+
+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 _FakeCommands:
+
+ def __init__(self, sandbox):
+ self._sandbox = sandbox
+
+ def run(self, command, timeout=None, background=False, cwd=None):
+ return self._sandbox.handle(command, background)
+
+
+class FakeSandbox:
+ """Stands in for an e2b sandbox: a filesystem plus a command channel.
+
+ 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:
+ # 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}))
+ 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 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):
+ 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 RemoteMsAgentToolEnvTest(unittest.TestCase):
+
+ def setUp(self):
+ self.tmp = tempfile.mkdtemp(prefix='envtest_')
+ 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.sandbox.requests[0]['calls'],
+ [{'tool_name': 'read_file', 'arguments': {'path': 'a.txt'}}])
+ self.assertEqual(result.observation, 'ran read_file')
+
+ 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(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_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('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
+ # 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):
+ # 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,
+ # which is indistinguishable from the task genuinely not being solved.
+ with self.assertRaises(ValueError):
+ self.env.resolve_tool('no_such_tool')
+
+ 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.
+ def responder(call):
+ if call['tool_name'] == 'code_executor---shell_executor':
+ return 'some output\n__TWINKLE_RC__:0'
+ return '__TWINKLE_RC__: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_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(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, env=env))
+ 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'}})))
+
+ 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."""
+
+ def setUp(self):
+ 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': {}}},
+ ]
+
+ 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(len(self.sandbox.requests), 1)
+ # 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
+
+ def explorer(trajectories, **kwargs):
+ return [{'messages': list(t['messages'])
+ + [{'role': 'assistant', 'content': '```python\nassert True\n```'}]}
+ for t in trajectories]
+
+ 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, envs=[self.env], 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.env.scripts, [])
+ # 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.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.env.scripts, ['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 = []
+
+ 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,
+ # 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,
+ **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,
+ envs=[FakeEnv()],
+ 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 environment.
+
+ A rack of one sandbox per slot is the whole point of running episodes in
+ 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
+ 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
+ # 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,
+ envs=envs,
+ 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 cleared twice, ran
+ # its own check twice, and every check saw the slot's own snapshot text.
+ from collections import Counter
+ 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(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.
+
+ 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, 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, env, **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],
+ envs=[env],
+ solver_rollouts=2,
+ keep_pass_band=(1, 2),
+ 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):
+ env = FakeEnv('input/a.csv 3\n')
+ ch = self._challenger(env)
+ kept = ch._filter_difficulty([self._task('#SETUP\nopen("a","w")')])
+ # 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."""
+
+ class NoSpaceEnv(FakeEnv):
+ """A workspace where the replay fails and a check would have passed."""
+
+ 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(env.scripts, [])
+ 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
+ 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 = []
+
+ 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]
+
+ 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,
+ envs=envs,
+ solver_rollouts=4,
+ keep_pass_band=(1, 4),
+ 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(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(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])
+
+
+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.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 = truncations.pop(0)
+ return [{'messages': list(t['messages']),
+ 'truncated': truncated,
+ 'stop_reason': 'length' if truncated else 'stop'}
+ for t in trajectories]
+
+ 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,
+ envs=[FakeEnv(exit_codes=[0 if passes else 1 for _, passes in attempt_flags])],
+ 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_pass_band=(1, 3))
+ 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_pass_band=(1, 3))
+ 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,
+ envs=[FakeEnv()],
+ keyword_store=self.store,
+ category_desc={'filesystem': 'files and directories'},
+ 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.keywords 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_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.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.
+ 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"]')
+ 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,
+ envs=[FakeEnv()],
+ keyword_store=self.store,
+ category_desc={category: 'some kind of work'},
+ 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.keywords._generate('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.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')
+ 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.keywords._AVOID_TOTAL
+ older = [f'old {i}' for i in range(200)]
+ fresh = [f'new {i}' for i in range(5)]
+ note = ch.keywords._avoid_note(older, fresh)
+ 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.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')
+ 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.
+
+ 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 recorder import Recorder
+
+ out = tempfile.mkdtemp(prefix='proposetraj_test_')
+ try:
+ 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(record['group_id'], 0)
+ self.assertEqual(record['reward'], 0.75)
+ self.assertEqual(record['side'], 'propose')
+
+
+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),
+ envs=[FakeEnv('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)
+
+
+# โโ 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()
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..7c1d36753
--- /dev/null
+++ b/tests/twinkle_agentic/test_harness.py
@@ -0,0 +1,248 @@
+# 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_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
+ 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
+ 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_multi_turn_rollout.py b/tests/twinkle_agentic/test_multi_turn_rollout.py
index 4f17d8282..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
@@ -156,6 +157,31 @@ 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 '')
+
# --- 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,21 +197,59 @@ 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
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,
@@ -198,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
@@ -216,9 +292,12 @@ def sample(self, pifs, sampling_params=None):
assert isinstance(pifs, list), (f'FakeSampler.sample expects a list, got {type(pifs).__name__}')
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'],
@@ -230,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."""
@@ -251,6 +334,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 +382,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 +453,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 +527,154 @@ 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_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],
+ 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
# =============================================================================
@@ -547,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
# =============================================================================
@@ -560,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()])
@@ -572,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'):
@@ -587,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()])
@@ -601,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'
@@ -622,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')])
@@ -670,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,
@@ -795,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')])
@@ -812,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')])
@@ -869,3 +1158,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_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_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..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'])
@@ -428,6 +446,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)
#
diff --git a/tests/utils/test_utils.py b/tests/utils/test_utils.py
index 641da8ad1..448d59909 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)]
@@ -358,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: