Merge 100a49bad2 into 393e387e9a
RileyChen23
GitHub
commit
c31f66525d
|
|
@ -1,6 +1,8 @@
|
|||
__pycache__/
|
||||
*.pyc
|
||||
.pytest_cache/
|
||||
.codex_pytest_tmp*/
|
||||
.DS_Store
|
||||
out
|
||||
website/
|
||||
docs-minimind/
|
||||
docs-minimind/
|
||||
|
|
|
|||
35
README.md
35
README.md
|
|
@ -1,28 +1,7 @@
|
|||
<div align="center">
|
||||
|
||||
|
||||
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
|
||||
|
||||
[](https://github.com/jingyaogong/minimind/stargazers)
|
||||
[](LICENSE)
|
||||
[](https://github.com/jingyaogong/minimind/commits/master)
|
||||
[](https://github.com/jingyaogong/minimind/pulls)
|
||||
[](https://huggingface.co/collections/jingyaogong/minimind-66caf8d999f5c7fa64f399e5)
|
||||
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
|
||||
|
||||
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
<h3>"大道至简"</h3>
|
||||
<h1>Lightweight LLM & Code Agent Lab</h1>
|
||||
<p>基于 MiniMind 的轻量模型训练、代码生成与执行反馈实验</p>
|
||||
<p><a href="./LICENSE">Apache-2.0 License</a></p>
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
|
|
@ -31,6 +10,8 @@
|
|||
|
||||
</div>
|
||||
|
||||
> 本仓库基于开源项目 [jingyaogong/minimind](https://github.com/jingyaogong/minimind) 进行复现与二次开发，保留原许可证、模型兼容性和上游致谢。新增的代码生成 Agent 支持代码提取、受控执行、隐藏测试反馈、pass@k 评估，以及面向 GRPO/CISPO 的可验证奖励。详见 [代码生成 Agent 文档](./docs/code-agent.md)。
|
||||
|
||||
* 此开源项目旨在完全从 0 开始，仅用 3 块钱成本与 2 小时训练时间，即可训练出规模约为 64M 的超小语言模型 MiniMind。
|
||||
* MiniMind 系列极其轻量，主线最小版本体积约为 GPT-3 的 $\frac{1}{2700}$，力求让普通个人 GPU 也能快速完成训练与复现。
|
||||
* 项目同时开源了大模型的极简结构与完整训练链路，覆盖 MoE、数据清洗、预训练（Pretrain）、监督微调（SFT）、LoRA、RLHF（DPO）、RLAIF（PPO / GRPO / CISPO）、Tool Use、Agentic RL、自适应思考与模型蒸馏等全过程代码。
|
||||
|
|
@ -786,8 +767,8 @@ LoRA 很适合处理“如何在尽量保留通用能力的前提下，让模型
|
|||
例2：自我认知数据
|
||||
|
||||
```jsonl
|
||||
{"conversations": [{"role": "user", "content": "你叫什么名字？"}, {"role": "assistant", "content": "您好，我名叫 MiniMind，是由 Jingyao Gong 开发的人工智能助手。我可以帮助你解决各种问题，包括翻译、推荐内容等等。"}]}
|
||||
{"conversations": [{"role": "user", "content": "你是谁"}, {"role": "assistant", "content": "您好，我是 MiniMind，由 Jingyao Gong 创建。我的主要职责是提供准确的回答和帮助。"}]}
|
||||
{"conversations": [{"role": "user", "content": "你叫什么名字？"}, {"role": "assistant", "content": "您好，我是一个开源的轻量级人工智能助手，可以帮助完成问答、翻译和内容整理等任务。"}]}
|
||||
{"conversations": [{"role": "user", "content": "你是谁"}, {"role": "assistant", "content": "您好，我是一个用于语言模型实验的开源助手，主要职责是提供准确的回答和帮助。"}]}
|
||||
```
|
||||
|
||||
此时可将【基模 + LoRA 权重】组合使用，从而在保留基础模型通用能力的同时，增强特定领域表现。训练完成后也可以通过 `eval_llm.py` 做简单测试。
|
||||
|
|
@ -802,7 +783,7 @@ python eval_llm.py --weight full_sft --lora_weight lora_medical
|
|||
python eval_llm.py --weight full_sft --lora_weight lora_identity
|
||||
|
||||
👶: 你和openai是什么关系？
|
||||
🤖️: 我是 MiniMind，一个由 Jingyao Gong 开发的人工智能助手。我通过自然语言处理和算法训练来与用户进行交互。
|
||||
🤖️: 我是一个开源的轻量级人工智能助手，通过自然语言处理模型与用户进行交互。
|
||||
```
|
||||
|
||||
PS：如果有更充足的数据，也可以直接做 `full_sft` 全参微调；不过这通常需要更谨慎地混合通用数据与领域数据，否则很容易因为过拟合垂域样本而损失模型原有的通用性。
|
||||
|
|
|
|||
35
README_en.md
35
README_en.md
|
|
@ -1,28 +1,7 @@
|
|||
<div align="center">
|
||||
|
||||
|
||||
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
|
||||
|
||||
[](https://github.com/jingyaogong/minimind/stargazers)
|
||||
[](LICENSE)
|
||||
[](https://github.com/jingyaogong/minimind/commits/master)
|
||||
[](https://github.com/jingyaogong/minimind/pulls)
|
||||
[](https://huggingface.co/collections/jingyaogong/minimind-66caf8d999f5c7fa64f399e5)
|
||||
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
|
||||
|
||||
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
<h3>"The Great Way is Simple"</h3>
|
||||
<h1>Lightweight LLM & Code Agent Lab</h1>
|
||||
<p>Lightweight model training, code generation, and execution-feedback experiments built on MiniMind</p>
|
||||
<p><a href="./LICENSE">Apache-2.0 License</a></p>
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
|
|
@ -31,6 +10,8 @@
|
|||
|
||||
</div>
|
||||
|
||||
> This repository reproduces and extends the open-source [jingyaogong/minimind](https://github.com/jingyaogong/minimind) project while preserving its license, model compatibility, and upstream attribution. The added code-generation agent provides code extraction, controlled execution, hidden-test feedback, pass@k evaluation, and verifiable rewards for GRPO/CISPO. See the [code agent documentation](./docs/code-agent.md).
|
||||
|
||||
* This open-source project aims to train MiniMind, an ultra-small language model with about 64M parameters, entirely from scratch with only about RMB 3 in cost and 2 hours of training time.
|
||||
* The MiniMind series is intentionally lightweight. The smallest model on the main branch is about $\frac{1}{2700}$ the size of GPT-3, making full training and reproduction feasible even on ordinary personal GPUs.
|
||||
* The project provides a minimalist model architecture and an end-to-end LLM training pipeline, covering MoE, data cleaning, pretraining, Supervised Fine-Tuning (SFT), LoRA, RLHF (DPO), RLAIF (PPO / GRPO / CISPO), Tool Use, Agentic RL, Adaptive Thinking, and Model Distillation.
|
||||
|
|
@ -786,8 +767,8 @@ Example 1: Vertical domain data
|
|||
Example 2: Self-awareness data
|
||||
|
||||
```jsonl
|
||||
{"conversations": [{"role": "user", "content": "你叫什么名字？"}, {"role": "assistant", "content": "您好，我名叫 MiniMind，是由 Jingyao Gong 开发的人工智能助手。我可以帮助你解决各种问题，包括翻译、推荐内容等等。"}]}
|
||||
{"conversations": [{"role": "user", "content": "你是谁"}, {"role": "assistant", "content": "您好，我是 MiniMind，由 Jingyao Gong 创建。我的主要职责是提供准确的回答和帮助。"}]}
|
||||
{"conversations": [{"role": "user", "content": "你叫什么名字？"}, {"role": "assistant", "content": "您好，我是一个开源的轻量级人工智能助手，可以帮助完成问答、翻译和内容整理等任务。"}]}
|
||||
{"conversations": [{"role": "user", "content": "你是谁"}, {"role": "assistant", "content": "您好，我是一个用于语言模型实验的开源助手，主要职责是提供准确的回答和帮助。"}]}
|
||||
```
|
||||
|
||||
At this point, the [base model + LoRA weights] can be used in combination, preserving the base model's general capabilities while enhancing specific domain performance. After training, simple tests can also be done via `eval_llm.py`.
|
||||
|
|
@ -802,7 +783,7 @@ python eval_llm.py --weight full_sft --lora_weight lora_medical
|
|||
python eval_llm.py --weight full_sft --lora_weight lora_identity
|
||||
|
||||
👶: 你和openai是什么关系？
|
||||
🤖️: 我是 MiniMind，一个由 Jingyao Gong 开发的人工智能助手。我通过自然语言处理和算法训练来与用户进行交互。
|
||||
🤖️: 我是一个开源的轻量级人工智能助手，通过自然语言处理模型与用户进行交互。
|
||||
```
|
||||
|
||||
PS: If you have sufficient data, you can also perform `full_sft` full-parameter fine-tuning directly. This usually requires careful mixing of general and domain-specific data; otherwise, the model may lose part of its general capability due to overfitting on vertical-domain samples.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,42 @@
|
|||
"""Utilities for execution-feedback code generation experiments."""
|
||||
|
||||
from .agent import AgentResult, AttemptRecord, ExecutionFeedbackAgent
|
||||
from .backends import OpenAICompatibleGenerator
|
||||
from .dataset import CodeRLDataset, collate_code_rl
|
||||
from .reward import CodeReward, score_code_batch, score_verification
|
||||
from .evaluation import (
|
||||
EvaluationReport,
|
||||
aggregate_reports,
|
||||
estimate_pass_at_k,
|
||||
evaluate_candidates,
|
||||
)
|
||||
from .verifier import (
|
||||
CodeTask,
|
||||
TestCase,
|
||||
VerificationResult,
|
||||
VerificationStatus,
|
||||
extract_python_code,
|
||||
verify_code,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AgentResult",
|
||||
"AttemptRecord",
|
||||
"CodeReward",
|
||||
"CodeRLDataset",
|
||||
"CodeTask",
|
||||
"EvaluationReport",
|
||||
"ExecutionFeedbackAgent",
|
||||
"OpenAICompatibleGenerator",
|
||||
"TestCase",
|
||||
"VerificationResult",
|
||||
"VerificationStatus",
|
||||
"aggregate_reports",
|
||||
"collate_code_rl",
|
||||
"estimate_pass_at_k",
|
||||
"evaluate_candidates",
|
||||
"extract_python_code",
|
||||
"score_verification",
|
||||
"score_code_batch",
|
||||
"verify_code",
|
||||
]
|
||||
|
|
@ -0,0 +1,144 @@
|
|||
"""A minimal generate-execute-feedback-repair loop for algorithmic code tasks."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import Callable
|
||||
|
||||
from .reward import CodeReward, score_verification
|
||||
from .verifier import (
|
||||
CodeTask,
|
||||
VerificationResult,
|
||||
VerificationStatus,
|
||||
extract_python_code,
|
||||
verify_code,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AttemptRecord:
|
||||
attempt: int
|
||||
response: str
|
||||
code: str
|
||||
verification: VerificationResult
|
||||
reward: CodeReward
|
||||
feedback_to_model: str
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
payload = asdict(self)
|
||||
payload["verification"] = self.verification.to_dict()
|
||||
payload["reward"] = asdict(self.reward)
|
||||
return payload
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AgentResult:
|
||||
task_id: str
|
||||
success: bool
|
||||
attempts: tuple[AttemptRecord, ...]
|
||||
|
||||
@property
|
||||
def final(self) -> AttemptRecord:
|
||||
return self.attempts[-1]
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"task_id": self.task_id,
|
||||
"success": self.success,
|
||||
"attempt_count": len(self.attempts),
|
||||
"attempts": [attempt.to_dict() for attempt in self.attempts],
|
||||
}
|
||||
|
||||
|
||||
def _feedback_for_model(result: VerificationResult, reveal_test_details: bool) -> str:
|
||||
if reveal_test_details:
|
||||
return result.feedback
|
||||
if result.status == VerificationStatus.WRONG_ANSWER:
|
||||
return (
|
||||
f"The solution passed {result.passed_tests}/{result.total_tests} tests. "
|
||||
"At least one hidden test failed. Re-check edge cases and algorithmic assumptions."
|
||||
)
|
||||
if result.status == VerificationStatus.RUNTIME_ERROR:
|
||||
last_line = next((line.strip() for line in reversed(result.stderr.splitlines()) if line.strip()), "")
|
||||
return f"The program raised a runtime error{': ' + last_line if last_line else ''}."
|
||||
return result.feedback
|
||||
|
||||
|
||||
class ExecutionFeedbackAgent:
|
||||
"""Iteratively ask a generator to repair code using verifier feedback.
|
||||
|
||||
``generate`` receives a plain prompt and returns a model response. This
|
||||
keeps the control loop independent from a specific inference backend.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
generate: Callable[[str], str],
|
||||
*,
|
||||
max_attempts: int = 3,
|
||||
timeout_seconds: float = 2.0,
|
||||
memory_mb: int = 256,
|
||||
reveal_test_details: bool = False,
|
||||
) -> None:
|
||||
if max_attempts < 1:
|
||||
raise ValueError("max_attempts must be at least 1")
|
||||
self.generate = generate
|
||||
self.max_attempts = max_attempts
|
||||
self.timeout_seconds = timeout_seconds
|
||||
self.memory_mb = memory_mb
|
||||
self.reveal_test_details = reveal_test_details
|
||||
|
||||
@staticmethod
|
||||
def _initial_prompt(task: CodeTask) -> str:
|
||||
return (
|
||||
"You are solving a Python algorithm task. Return exactly one ```python``` code block.\n"
|
||||
f"Required entry point: {task.entry_point}\n"
|
||||
f"Task:\n{task.prompt}\n"
|
||||
"Do not read files, access the network, spawn processes, or call unsafe builtins."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _repair_prompt(task: CodeTask, previous_code: str, feedback: str, attempt: int) -> str:
|
||||
return (
|
||||
"Repair the previous Python solution using the execution feedback. "
|
||||
"Return the complete replacement in exactly one ```python``` code block.\n"
|
||||
f"Required entry point: {task.entry_point}\n"
|
||||
f"Task:\n{task.prompt}\n"
|
||||
f"Previous solution (attempt {attempt - 1}):\n```python\n{previous_code}\n```\n"
|
||||
f"Execution feedback:\n{feedback}"
|
||||
)
|
||||
|
||||
def run(self, task: CodeTask) -> AgentResult:
|
||||
attempts: list[AttemptRecord] = []
|
||||
prompt = self._initial_prompt(task)
|
||||
previous_code = ""
|
||||
for attempt_number in range(1, self.max_attempts + 1):
|
||||
response = self.generate(prompt)
|
||||
code = extract_python_code(response)
|
||||
verification = verify_code(
|
||||
code,
|
||||
task,
|
||||
timeout_seconds=self.timeout_seconds,
|
||||
memory_mb=self.memory_mb,
|
||||
)
|
||||
reward = score_verification(verification)
|
||||
feedback = _feedback_for_model(verification, self.reveal_test_details)
|
||||
attempts.append(
|
||||
AttemptRecord(
|
||||
attempt=attempt_number,
|
||||
response=response,
|
||||
code=code,
|
||||
verification=verification,
|
||||
reward=reward,
|
||||
feedback_to_model=feedback,
|
||||
)
|
||||
)
|
||||
if verification.status == VerificationStatus.PASSED:
|
||||
break
|
||||
previous_code = code
|
||||
prompt = self._repair_prompt(task, previous_code, feedback, attempt_number + 1)
|
||||
return AgentResult(
|
||||
task_id=task.task_id,
|
||||
success=attempts[-1].verification.status == VerificationStatus.PASSED,
|
||||
attempts=tuple(attempts),
|
||||
)
|
||||
|
|
@ -0,0 +1,57 @@
|
|||
"""Inference adapters for the execution-feedback agent."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
class OpenAICompatibleGenerator:
|
||||
"""Callable adapter for MiniMind, vLLM, SGLang, or other OpenAI-style APIs."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
base_url: str = "http://localhost:8998/v1",
|
||||
api_key: str = "not-needed",
|
||||
model: str = "minimind",
|
||||
temperature: float = 0.7,
|
||||
max_tokens: int = 1024,
|
||||
open_thinking: bool = False,
|
||||
seed: int | None = None,
|
||||
client: Any = None,
|
||||
) -> None:
|
||||
if client is None:
|
||||
try:
|
||||
from openai import OpenAI
|
||||
except ImportError as exc:
|
||||
raise RuntimeError("Install the project 'openai' dependency to use this backend.") from exc
|
||||
client = OpenAI(api_key=api_key, base_url=base_url)
|
||||
self.client = client
|
||||
self.model = model
|
||||
self.temperature = temperature
|
||||
self.max_tokens = max_tokens
|
||||
self.open_thinking = open_thinking
|
||||
self.seed = seed
|
||||
self._request_index = 0
|
||||
|
||||
def __call__(self, prompt: str) -> str:
|
||||
request = {
|
||||
"model": self.model,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"temperature": self.temperature,
|
||||
"max_tokens": self.max_tokens,
|
||||
"stream": False,
|
||||
"extra_body": {"open_thinking": self.open_thinking},
|
||||
}
|
||||
if self.seed is not None:
|
||||
request["seed"] = self.seed + self._request_index
|
||||
self._request_index += 1
|
||||
response = self.client.chat.completions.create(
|
||||
**request,
|
||||
)
|
||||
if not response.choices or response.choices[0].message is None:
|
||||
raise RuntimeError("Model API returned no completion choice.")
|
||||
content = response.choices[0].message.content
|
||||
if not content:
|
||||
raise RuntimeError("Model API returned an empty completion.")
|
||||
return content
|
||||
|
|
@ -0,0 +1,58 @@
|
|||
"""Dataset adapter for execution-verifiable GRPO code tasks."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import random
|
||||
from pathlib import Path
|
||||
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
from .verifier import CodeTask
|
||||
|
||||
|
||||
class CodeRLDataset(Dataset):
|
||||
def __init__(self, jsonl_path, tokenizer, thinking_ratio: float = 0.5):
|
||||
if not 0.0 <= thinking_ratio <= 1.0:
|
||||
raise ValueError("thinking_ratio must be between 0 and 1")
|
||||
self.tokenizer = tokenizer
|
||||
self.thinking_ratio = thinking_ratio
|
||||
self.tasks: list[CodeTask] = []
|
||||
with Path(jsonl_path).open("r", encoding="utf-8") as handle:
|
||||
for line_number, line in enumerate(handle, start=1):
|
||||
if not line.strip():
|
||||
continue
|
||||
try:
|
||||
self.tasks.append(CodeTask.from_dict(json.loads(line)))
|
||||
except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
|
||||
raise ValueError(f"Invalid code task at {jsonl_path}:{line_number}: {exc}") from exc
|
||||
if not self.tasks:
|
||||
raise ValueError("code RL dataset is empty")
|
||||
|
||||
def __len__(self):
|
||||
return len(self.tasks)
|
||||
|
||||
def __getitem__(self, index):
|
||||
task = self.tasks[index]
|
||||
message = {
|
||||
"role": "user",
|
||||
"content": (
|
||||
"Solve the Python algorithm task below. Return one complete ```python``` code block.\n"
|
||||
f"Required entry point: {task.entry_point}\n"
|
||||
f"Task:\n{task.prompt}"
|
||||
),
|
||||
}
|
||||
prompt = self.tokenizer.apply_chat_template(
|
||||
[message],
|
||||
tokenize=False,
|
||||
open_thinking=random.random() < self.thinking_ratio,
|
||||
add_generation_prompt=True,
|
||||
)
|
||||
return {"prompt": prompt, "task": task}
|
||||
|
||||
|
||||
def collate_code_rl(samples):
|
||||
return {
|
||||
"prompt": [sample["prompt"] for sample in samples],
|
||||
"tasks": [sample["task"] for sample in samples],
|
||||
}
|
||||
|
|
@ -0,0 +1,116 @@
|
|||
"""Evaluation metrics for execution-verified code generation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from collections import Counter
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import Any, Iterable, Sequence
|
||||
|
||||
from .verifier import CodeTask, VerificationResult, VerificationStatus, verify_code
|
||||
|
||||
|
||||
def estimate_pass_at_k(num_samples: int, num_correct: int, k: int) -> float:
|
||||
"""Unbiased pass@k estimator used by common code-generation benchmarks."""
|
||||
if num_samples <= 0:
|
||||
raise ValueError("num_samples must be positive")
|
||||
if not 0 <= num_correct <= num_samples:
|
||||
raise ValueError("num_correct must be between 0 and num_samples")
|
||||
if not 1 <= k <= num_samples:
|
||||
raise ValueError("k must be between 1 and num_samples")
|
||||
if num_samples - num_correct < k:
|
||||
return 1.0
|
||||
return 1.0 - math.comb(num_samples - num_correct, k) / math.comb(num_samples, k)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class EvaluationReport:
|
||||
task_id: str
|
||||
candidate_count: int
|
||||
correct_count: int
|
||||
compile_rate: float
|
||||
execution_rate: float
|
||||
test_pass_rate: float
|
||||
pass_at_k: dict[str, float]
|
||||
status_counts: dict[str, int]
|
||||
results: tuple[VerificationResult, ...]
|
||||
|
||||
def to_dict(self, *, include_results: bool = True) -> dict[str, Any]:
|
||||
payload = asdict(self)
|
||||
if include_results:
|
||||
payload["results"] = [result.to_dict() for result in self.results]
|
||||
else:
|
||||
payload.pop("results", None)
|
||||
return payload
|
||||
|
||||
|
||||
_COMPILED_STATUSES = {
|
||||
VerificationStatus.PASSED,
|
||||
VerificationStatus.RUNTIME_ERROR,
|
||||
VerificationStatus.TIMEOUT,
|
||||
VerificationStatus.WRONG_ANSWER,
|
||||
}
|
||||
_EXECUTED_STATUSES = {
|
||||
VerificationStatus.PASSED,
|
||||
VerificationStatus.RUNTIME_ERROR,
|
||||
VerificationStatus.WRONG_ANSWER,
|
||||
}
|
||||
|
||||
|
||||
def evaluate_candidates(
|
||||
task: CodeTask,
|
||||
candidates: Sequence[str],
|
||||
*,
|
||||
ks: Iterable[int] = (1, 5, 10),
|
||||
timeout_seconds: float = 2.0,
|
||||
memory_mb: int = 256,
|
||||
) -> EvaluationReport:
|
||||
if not candidates:
|
||||
raise ValueError("at least one candidate is required")
|
||||
results = tuple(
|
||||
verify_code(candidate, task, timeout_seconds=timeout_seconds, memory_mb=memory_mb)
|
||||
for candidate in candidates
|
||||
)
|
||||
count = len(results)
|
||||
correct = sum(result.status == VerificationStatus.PASSED for result in results)
|
||||
status_counts = Counter(result.status.value for result in results)
|
||||
total_tests = sum(result.total_tests for result in results)
|
||||
passed_tests = sum(result.passed_tests for result in results)
|
||||
requested_ks = sorted(set(int(k) for k in ks if 1 <= int(k) <= count))
|
||||
pass_at_k = {f"pass@{k}": estimate_pass_at_k(count, correct, k) for k in requested_ks}
|
||||
return EvaluationReport(
|
||||
task_id=task.task_id,
|
||||
candidate_count=count,
|
||||
correct_count=correct,
|
||||
compile_rate=sum(result.status in _COMPILED_STATUSES for result in results) / count,
|
||||
execution_rate=sum(result.status in _EXECUTED_STATUSES for result in results) / count,
|
||||
test_pass_rate=passed_tests / total_tests if total_tests else 0.0,
|
||||
pass_at_k=pass_at_k,
|
||||
status_counts=dict(sorted(status_counts.items())),
|
||||
results=results,
|
||||
)
|
||||
|
||||
|
||||
def aggregate_reports(reports: Sequence[EvaluationReport]) -> dict[str, Any]:
|
||||
if not reports:
|
||||
raise ValueError("at least one evaluation report is required")
|
||||
status_counts: Counter[str] = Counter()
|
||||
for report in reports:
|
||||
status_counts.update(report.status_counts)
|
||||
metric_names = sorted({name for report in reports for name in report.pass_at_k})
|
||||
pass_at_k = {
|
||||
name: sum(report.pass_at_k[name] for report in reports if name in report.pass_at_k)
|
||||
/ sum(name in report.pass_at_k for report in reports)
|
||||
for name in metric_names
|
||||
}
|
||||
candidate_count = sum(report.candidate_count for report in reports)
|
||||
return {
|
||||
"task_count": len(reports),
|
||||
"candidate_count": candidate_count,
|
||||
"correct_count": sum(report.correct_count for report in reports),
|
||||
"compile_rate": sum(report.compile_rate * report.candidate_count for report in reports) / candidate_count,
|
||||
"execution_rate": sum(report.execution_rate * report.candidate_count for report in reports) / candidate_count,
|
||||
"test_pass_rate": sum(report.test_pass_rate * report.candidate_count for report in reports) / candidate_count,
|
||||
"pass_at_k": pass_at_k,
|
||||
"status_counts": dict(sorted(status_counts.items())),
|
||||
}
|
||||
|
|
@ -0,0 +1,65 @@
|
|||
"""Verifiable rewards derived from deterministic code execution."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Sequence
|
||||
|
||||
from .verifier import CodeTask, VerificationResult, VerificationStatus, verify_code
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CodeReward:
|
||||
total: float
|
||||
correctness: float
|
||||
execution: float
|
||||
format: float
|
||||
reason: str
|
||||
|
||||
|
||||
def score_verification(result: VerificationResult) -> CodeReward:
|
||||
"""Map verifier outcomes to a transparent reward in [-1, 1]."""
|
||||
status = result.status
|
||||
if status == VerificationStatus.PASSED:
|
||||
return CodeReward(1.0, 0.8, 0.1, 0.1, "all tests passed")
|
||||
if status == VerificationStatus.WRONG_ANSWER:
|
||||
correctness = 0.8 * result.pass_rate
|
||||
total = -0.2 + correctness + 0.1 + 0.1
|
||||
return CodeReward(total, correctness, 0.1, 0.1, "partial test credit")
|
||||
if status == VerificationStatus.RUNTIME_ERROR:
|
||||
return CodeReward(-0.4, 0.0, -0.5, 0.1, "runtime error")
|
||||
if status == VerificationStatus.TIMEOUT:
|
||||
return CodeReward(-0.6, 0.0, -0.7, 0.1, "execution timeout")
|
||||
if status == VerificationStatus.SYNTAX_ERROR:
|
||||
return CodeReward(-0.8, 0.0, -0.8, 0.0, "syntax error")
|
||||
if status in {VerificationStatus.NO_CODE, VerificationStatus.POLICY_VIOLATION}:
|
||||
return CodeReward(-1.0, 0.0, 0.0, -1.0, status.value)
|
||||
return CodeReward(-1.0, 0.0, -1.0, 0.0, "verifier failure")
|
||||
|
||||
|
||||
def score_code_batch(
|
||||
tasks: Sequence[CodeTask],
|
||||
responses: Sequence[str],
|
||||
*,
|
||||
num_generations: int,
|
||||
timeout_seconds: float = 2.0,
|
||||
memory_mb: int = 256,
|
||||
) -> list[CodeReward]:
|
||||
"""Score grouped generations in task-major order for GRPO/CISPO."""
|
||||
if num_generations < 1:
|
||||
raise ValueError("num_generations must be at least 1")
|
||||
expected = len(tasks) * num_generations
|
||||
if len(responses) != expected:
|
||||
raise ValueError(f"expected {expected} responses for {len(tasks)} tasks, got {len(responses)}")
|
||||
rewards: list[CodeReward] = []
|
||||
for task_index, task in enumerate(tasks):
|
||||
start = task_index * num_generations
|
||||
for response in responses[start:start + num_generations]:
|
||||
result = verify_code(
|
||||
response,
|
||||
task,
|
||||
timeout_seconds=timeout_seconds,
|
||||
memory_mb=memory_mb,
|
||||
)
|
||||
rewards.append(score_verification(result))
|
||||
return rewards
|
||||
|
|
@ -0,0 +1,351 @@
|
|||
"""Deterministic verifier for function-style Python code generation tasks.
|
||||
|
||||
The local subprocess backend is intended for controlled experiments. It uses
|
||||
AST policy checks, a temporary working directory, an isolated Python process,
|
||||
and a hard timeout, but it is not a production security boundary. Run code
|
||||
from untrusted users inside a dedicated container or micro-VM.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import textwrap
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
class VerificationStatus(str, Enum):
|
||||
PASSED = "passed"
|
||||
NO_CODE = "no_code"
|
||||
SYNTAX_ERROR = "syntax_error"
|
||||
POLICY_VIOLATION = "policy_violation"
|
||||
RUNTIME_ERROR = "runtime_error"
|
||||
TIMEOUT = "timeout"
|
||||
WRONG_ANSWER = "wrong_answer"
|
||||
INTERNAL_ERROR = "internal_error"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TestCase:
|
||||
args: tuple[Any, ...] = ()
|
||||
kwargs: dict[str, Any] = field(default_factory=dict)
|
||||
expected: Any = None
|
||||
|
||||
def to_json(self) -> dict[str, Any]:
|
||||
return {"args": list(self.args), "kwargs": self.kwargs, "expected": self.expected}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, payload: dict[str, Any]) -> "TestCase":
|
||||
return cls(
|
||||
args=tuple(payload.get("args", ())),
|
||||
kwargs=dict(payload.get("kwargs", {})),
|
||||
expected=payload.get("expected"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CodeTask:
|
||||
task_id: str
|
||||
prompt: str
|
||||
entry_point: str
|
||||
tests: tuple[TestCase, ...]
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if not self.entry_point.isidentifier():
|
||||
raise ValueError(f"entry_point must be a Python identifier: {self.entry_point!r}")
|
||||
if not self.tests:
|
||||
raise ValueError("a code task must contain at least one test")
|
||||
# Fail early instead of discovering non-serializable tests in a worker.
|
||||
json.dumps([case.to_json() for case in self.tests], ensure_ascii=False)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"task_id": self.task_id,
|
||||
"prompt": self.prompt,
|
||||
"entry_point": self.entry_point,
|
||||
"tests": [case.to_json() for case in self.tests],
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, payload: dict[str, Any]) -> "CodeTask":
|
||||
return cls(
|
||||
task_id=str(payload["task_id"]),
|
||||
prompt=str(payload["prompt"]),
|
||||
entry_point=str(payload["entry_point"]),
|
||||
tests=tuple(TestCase.from_dict(case) for case in payload["tests"]),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class VerificationResult:
|
||||
status: VerificationStatus
|
||||
passed_tests: int
|
||||
total_tests: int
|
||||
duration_ms: float
|
||||
feedback: str
|
||||
stdout: str = ""
|
||||
stderr: str = ""
|
||||
|
||||
@property
|
||||
def pass_rate(self) -> float:
|
||||
return self.passed_tests / self.total_tests if self.total_tests else 0.0
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
payload = asdict(self)
|
||||
payload["status"] = self.status.value
|
||||
payload["pass_rate"] = self.pass_rate
|
||||
return payload
|
||||
|
||||
|
||||
_FENCED_CODE = re.compile(r"```(?:python|py)?\s*\n(.*?)```", re.IGNORECASE | re.DOTALL)
|
||||
_THINK_BLOCK = re.compile(r"<think>.*?</think>", re.IGNORECASE | re.DOTALL)
|
||||
|
||||
_ALLOWED_IMPORT_ROOTS = {
|
||||
"bisect",
|
||||
"collections",
|
||||
"functools",
|
||||
"heapq",
|
||||
"itertools",
|
||||
"math",
|
||||
"operator",
|
||||
"re",
|
||||
"string",
|
||||
"typing",
|
||||
}
|
||||
_BLOCKED_CALLS = {"breakpoint", "compile", "eval", "exec", "input", "open", "__import__"}
|
||||
_BLOCKED_NAMES = {"os", "pathlib", "requests", "shutil", "socket", "subprocess", "sys"}
|
||||
|
||||
|
||||
def extract_python_code(response: str) -> str:
|
||||
"""Extract the first fenced Python block, falling back to plain text."""
|
||||
cleaned = _THINK_BLOCK.sub("", response or "").strip()
|
||||
matches = _FENCED_CODE.findall(cleaned)
|
||||
if matches:
|
||||
return textwrap.dedent(matches[0]).strip()
|
||||
return textwrap.dedent(cleaned).strip()
|
||||
|
||||
|
||||
def _policy_errors(tree: ast.AST) -> list[str]:
|
||||
errors: list[str] = []
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.Import):
|
||||
for alias in node.names:
|
||||
root = alias.name.split(".", 1)[0]
|
||||
if root not in _ALLOWED_IMPORT_ROOTS:
|
||||
errors.append(f"import {alias.name!r} is not allowed")
|
||||
elif isinstance(node, ast.ImportFrom):
|
||||
root = (node.module or "").split(".", 1)[0]
|
||||
if root not in _ALLOWED_IMPORT_ROOTS:
|
||||
errors.append(f"import from {node.module!r} is not allowed")
|
||||
elif isinstance(node, ast.Call) and isinstance(node.func, ast.Name):
|
||||
if node.func.id in _BLOCKED_CALLS:
|
||||
errors.append(f"call to {node.func.id!r} is not allowed")
|
||||
elif isinstance(node, ast.Name) and node.id in _BLOCKED_NAMES:
|
||||
errors.append(f"name {node.id!r} is not allowed")
|
||||
elif isinstance(node, ast.Attribute) and node.attr.startswith("__"):
|
||||
errors.append(f"dunder attribute {node.attr!r} is not allowed")
|
||||
return sorted(set(errors))
|
||||
|
||||
|
||||
def _limit_resources(timeout_seconds: float, memory_mb: int):
|
||||
if os.name == "nt":
|
||||
return None
|
||||
|
||||
def apply_limits() -> None:
|
||||
import resource
|
||||
|
||||
cpu_seconds = max(1, math.ceil(timeout_seconds))
|
||||
memory_bytes = memory_mb * 1024 * 1024
|
||||
resource.setrlimit(resource.RLIMIT_CPU, (cpu_seconds, cpu_seconds))
|
||||
resource.setrlimit(resource.RLIMIT_AS, (memory_bytes, memory_bytes))
|
||||
resource.setrlimit(resource.RLIMIT_FSIZE, (1024 * 1024, 1024 * 1024))
|
||||
|
||||
return apply_limits
|
||||
|
||||
|
||||
_RUNNER_SOURCE = r'''
|
||||
import importlib.util
|
||||
import json
|
||||
import math
|
||||
import traceback
|
||||
|
||||
|
||||
def equivalent(actual, expected):
|
||||
if isinstance(actual, float) or isinstance(expected, float):
|
||||
try:
|
||||
return math.isclose(float(actual), float(expected), rel_tol=1e-7, abs_tol=1e-9)
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
if isinstance(actual, tuple):
|
||||
actual = list(actual)
|
||||
if isinstance(expected, tuple):
|
||||
expected = list(expected)
|
||||
return actual == expected
|
||||
|
||||
|
||||
def main():
|
||||
spec = importlib.util.spec_from_file_location("candidate_solution", "solution.py")
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
with open("cases.json", "r", encoding="utf-8") as handle:
|
||||
payload = json.load(handle)
|
||||
function = getattr(module, payload["entry_point"])
|
||||
passed = 0
|
||||
failure = None
|
||||
for index, case in enumerate(payload["tests"]):
|
||||
actual = function(*case["args"], **case["kwargs"])
|
||||
if equivalent(actual, case["expected"]):
|
||||
passed += 1
|
||||
elif failure is None:
|
||||
failure = {
|
||||
"index": index,
|
||||
"expected": repr(case["expected"]),
|
||||
"actual": repr(actual),
|
||||
}
|
||||
print("__MINIMIND_RESULT__=" + json.dumps({"passed": passed, "failure": failure}))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
main()
|
||||
except Exception:
|
||||
print("__MINIMIND_EXCEPTION__=" + traceback.format_exc())
|
||||
raise
|
||||
'''
|
||||
|
||||
|
||||
def _clip_output(value: str, limit: int = 4000) -> str:
|
||||
return value if len(value) <= limit else value[-limit:]
|
||||
|
||||
|
||||
def verify_code(
|
||||
response_or_code: str,
|
||||
task: CodeTask,
|
||||
*,
|
||||
timeout_seconds: float = 2.0,
|
||||
memory_mb: int = 256,
|
||||
) -> VerificationResult:
|
||||
"""Run a candidate function against deterministic JSON-serializable tests."""
|
||||
started = time.perf_counter()
|
||||
code = extract_python_code(response_or_code)
|
||||
if not code:
|
||||
return VerificationResult(
|
||||
VerificationStatus.NO_CODE, 0, len(task.tests), 0.0, "No Python code was found."
|
||||
)
|
||||
|
||||
try:
|
||||
tree = ast.parse(code)
|
||||
except SyntaxError as exc:
|
||||
duration = (time.perf_counter() - started) * 1000
|
||||
return VerificationResult(
|
||||
VerificationStatus.SYNTAX_ERROR,
|
||||
0,
|
||||
len(task.tests),
|
||||
duration,
|
||||
f"SyntaxError on line {exc.lineno}: {exc.msg}",
|
||||
)
|
||||
|
||||
errors = _policy_errors(tree)
|
||||
if errors:
|
||||
duration = (time.perf_counter() - started) * 1000
|
||||
return VerificationResult(
|
||||
VerificationStatus.POLICY_VIOLATION,
|
||||
0,
|
||||
len(task.tests),
|
||||
duration,
|
||||
"; ".join(errors),
|
||||
)
|
||||
|
||||
payload = {"entry_point": task.entry_point, "tests": [case.to_json() for case in task.tests]}
|
||||
env = {"PYTHONHASHSEED": "0", "PYTHONIOENCODING": "utf-8"}
|
||||
creationflags = subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0
|
||||
|
||||
try:
|
||||
with tempfile.TemporaryDirectory(prefix="minimind-code-") as temp_dir:
|
||||
root = Path(temp_dir)
|
||||
(root / "solution.py").write_text(code, encoding="utf-8")
|
||||
(root / "cases.json").write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8")
|
||||
(root / "runner.py").write_text(_RUNNER_SOURCE, encoding="utf-8")
|
||||
completed = subprocess.run(
|
||||
[sys.executable, "-I", "-S", "runner.py"],
|
||||
cwd=root,
|
||||
env=env,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
encoding="utf-8",
|
||||
errors="replace",
|
||||
timeout=timeout_seconds,
|
||||
check=False,
|
||||
creationflags=creationflags,
|
||||
preexec_fn=_limit_resources(timeout_seconds, memory_mb),
|
||||
)
|
||||
except subprocess.TimeoutExpired as exc:
|
||||
duration = (time.perf_counter() - started) * 1000
|
||||
return VerificationResult(
|
||||
VerificationStatus.TIMEOUT,
|
||||
0,
|
||||
len(task.tests),
|
||||
duration,
|
||||
f"Execution exceeded the {timeout_seconds:.2f}s timeout.",
|
||||
_clip_output(exc.stdout or ""),
|
||||
_clip_output(exc.stderr or ""),
|
||||
)
|
||||
except Exception as exc:
|
||||
duration = (time.perf_counter() - started) * 1000
|
||||
return VerificationResult(
|
||||
VerificationStatus.INTERNAL_ERROR,
|
||||
0,
|
||||
len(task.tests),
|
||||
duration,
|
||||
f"Verifier failure: {type(exc).__name__}: {exc}",
|
||||
)
|
||||
|
||||
duration = (time.perf_counter() - started) * 1000
|
||||
stdout, stderr = _clip_output(completed.stdout), _clip_output(completed.stderr)
|
||||
marker = "__MINIMIND_RESULT__="
|
||||
result_line = next((line for line in reversed(stdout.splitlines()) if line.startswith(marker)), None)
|
||||
if completed.returncode != 0 or result_line is None:
|
||||
exception_line = next(
|
||||
(line for line in reversed(stdout.splitlines()) if line.startswith("__MINIMIND_EXCEPTION__=")),
|
||||
None,
|
||||
)
|
||||
feedback = exception_line.split("=", 1)[1] if exception_line else (stderr.strip() or "Candidate raised an exception.")
|
||||
return VerificationResult(
|
||||
VerificationStatus.RUNTIME_ERROR, 0, len(task.tests), duration, feedback, stdout, stderr
|
||||
)
|
||||
|
||||
try:
|
||||
worker_result = json.loads(result_line[len(marker):])
|
||||
passed = int(worker_result["passed"])
|
||||
except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
|
||||
return VerificationResult(
|
||||
VerificationStatus.INTERNAL_ERROR,
|
||||
0,
|
||||
len(task.tests),
|
||||
duration,
|
||||
f"Invalid verifier output: {exc}",
|
||||
stdout,
|
||||
stderr,
|
||||
)
|
||||
|
||||
failure = worker_result.get("failure")
|
||||
if passed == len(task.tests):
|
||||
status = VerificationStatus.PASSED
|
||||
feedback = f"Passed all {passed}/{len(task.tests)} tests."
|
||||
else:
|
||||
status = VerificationStatus.WRONG_ANSWER
|
||||
feedback = (
|
||||
f"Passed {passed}/{len(task.tests)} tests. First failure at test {failure['index']}: "
|
||||
f"expected {failure['expected']}, got {failure['actual']}."
|
||||
)
|
||||
return VerificationResult(status, passed, len(task.tests), duration, feedback, stdout, stderr)
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
{"task_id":"reverse-string","prompt":"Implement reverse_string(text), returning the characters in reverse order.","entry_point":"reverse_string","tests":[{"args":["hello"],"expected":"olleh"},{"args":[""],"expected":""},{"args":["a b"],"expected":"b a"}]}
|
||||
{"task_id":"valid-parentheses","prompt":"Implement valid_parentheses(text). The input contains only ()[]{}. Return whether every bracket is correctly matched and nested.","entry_point":"valid_parentheses","tests":[{"args":["()[]{}"],"expected":true},{"args":["([{}])"],"expected":true},{"args":["(]"],"expected":false},{"args":["([)]"],"expected":false}]}
|
||||
{"task_id":"binary-search","prompt":"Implement binary_search(values, target) for an ascending list of distinct integers. Return the target index or -1 when absent.","entry_point":"binary_search","tests":[{"args":[[1,3,5,7],5],"expected":2},{"args":[[1,3,5,7],2],"expected":-1},{"args":[[],4],"expected":-1}]}
|
||||
{"task_id":"max-subarray","prompt":"Implement max_subarray(values), returning the maximum sum of a non-empty contiguous subarray.","entry_point":"max_subarray","tests":[{"args":[[-2,1,-3,4,-1,2,1,-5,4]],"expected":6},{"args":[[-5,-2,-8]],"expected":-2},{"args":[[7]],"expected":7}]}
|
||||
|
|
@ -0,0 +1,143 @@
|
|||
# Code Generation Agent
|
||||
|
||||
This extension adds an execution-feedback loop for Python algorithm generation while keeping the original MiniMind model and checkpoint format unchanged.
|
||||
|
||||
## Architecture
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A[Algorithm task] --> B[MiniMind or OpenAI-compatible model]
|
||||
B --> C[Python code extraction]
|
||||
C --> D[AST policy validation]
|
||||
D --> E[Isolated subprocess]
|
||||
E --> F[Deterministic tests]
|
||||
F --> G[Structured feedback]
|
||||
G -->|repair prompt| B
|
||||
F --> H[Verifiable reward]
|
||||
H --> I[GRPO or CISPO]
|
||||
F --> J[pass@k and error metrics]
|
||||
```
|
||||
|
||||
The implementation is split into small, testable modules:
|
||||
|
||||
- `code_agent/verifier.py`: code extraction, AST policy checks, subprocess execution, timeout and structured results.
|
||||
- `code_agent/reward.py`: transparent reward components and grouped GRPO scoring.
|
||||
- `code_agent/agent.py`: generate-execute-feedback-repair loop with hidden-test-safe feedback.
|
||||
- `code_agent/evaluation.py`: compile rate, execution rate, test pass rate, error counts and unbiased pass@k.
|
||||
- `code_agent/dataset.py`: JSONL code tasks adapted to the MiniMind chat template.
|
||||
- `code_agent/backends.py`: OpenAI-compatible inference adapter for MiniMind, vLLM and SGLang.
|
||||
|
||||
## Task format
|
||||
|
||||
Each JSONL record contains a prompt, a required function name, and JSON-serializable tests:
|
||||
|
||||
```json
|
||||
{"task_id":"add-two","prompt":"Implement add(a, b).","entry_point":"add","tests":[{"args":[1,2],"expected":3}]}
|
||||
```
|
||||
|
||||
`dataset/code_rl_mini.jsonl` contains four smoke-test tasks. It is deliberately too small for a meaningful training claim; use it to validate the pipeline before preparing a larger train/validation/test split.
|
||||
|
||||
For a standard public benchmark, convert the compatible subset of MBPP-sanitized while preserving its official train/validation/test split:
|
||||
|
||||
```bash
|
||||
python scripts/prepare_mbpp.py --output-dir out/mbpp_sanitized
|
||||
```
|
||||
|
||||
The converter accepts direct, JSON-serializable function assertions, validates every converted task against the published reference solution, and reports all skipped records. It also emits `train_sft.jsonl` from the official training references, enabling an SFT → execution-reward GRPO ablation without leaking validation or test solutions. Generated benchmark files remain under the ignored `out/` directory.
|
||||
|
||||
## Evaluate saved candidates
|
||||
|
||||
Generate multiple candidates from a running OpenAI-compatible model service with deterministic per-request seeds:
|
||||
|
||||
```bash
|
||||
python scripts/generate_codegen.py \
|
||||
--tasks dataset/code_rl_mini.jsonl \
|
||||
--samples-per-task 10 \
|
||||
--seed 42 \
|
||||
--output out/code_predictions.jsonl
|
||||
```
|
||||
|
||||
Each completed task is flushed to JSONL immediately. If a long run is interrupted, rerun the same command with `--resume`; task-specific seed ranges keep completed and resumed outputs reproducible.
|
||||
|
||||
Then execute and aggregate them:
|
||||
|
||||
```bash
|
||||
python scripts/eval_codegen.py \
|
||||
--tasks dataset/code_rl_mini.jsonl \
|
||||
--predictions out/code_predictions.jsonl \
|
||||
--ks 1 5 10
|
||||
```
|
||||
|
||||
The report separates:
|
||||
|
||||
- `compile_rate`: candidate passed syntax and policy validation;
|
||||
- `execution_rate`: candidate completed without a timeout or runtime exception;
|
||||
- `test_pass_rate`: fraction of deterministic tests passed;
|
||||
- `pass@k`: probability that at least one of `k` sampled candidates is correct;
|
||||
- `status_counts`: syntax, policy, timeout, runtime and wrong-answer failure modes.
|
||||
|
||||
## Run the repair agent
|
||||
|
||||
Start the repository's OpenAI-compatible service, then run:
|
||||
|
||||
```bash
|
||||
python scripts/run_code_agent.py \
|
||||
--tasks dataset/code_rl_mini.jsonl \
|
||||
--base-url http://localhost:8998/v1 \
|
||||
--model minimind \
|
||||
--max-attempts 3 \
|
||||
--seed 42 \
|
||||
--output out/code_agent.json \
|
||||
--resume
|
||||
```
|
||||
|
||||
Each task receives a stable seed range, so the initial attempt remains comparable across single-pass and repair runs. With `--resume`, the output is atomically checkpointed after every task and completed tasks are skipped after an interruption. Hidden-test details are not sent back to the model by default. Add `--reveal-test-details` only for public development tests.
|
||||
|
||||
## Train with verifiable rewards
|
||||
|
||||
From the repository root:
|
||||
|
||||
```bash
|
||||
cd trainer
|
||||
python train_grpo.py \
|
||||
--reward_type code \
|
||||
--data_path ../dataset/code_rl_mini.jsonl \
|
||||
--num_generations 6 \
|
||||
--loss_type grpo \
|
||||
--from_weight full_sft
|
||||
```
|
||||
|
||||
The training log includes `degenerate_group_rate`. A group whose candidates all receive the same reward has zero normalized advantage and contributes no useful policy-gradient signal. This metric makes the failure mode visible before introducing strategies such as dynamic sampling or harder tasks.
|
||||
|
||||
## Reward design
|
||||
|
||||
The default reward is intentionally transparent and bounded to `[-1, 1]`:
|
||||
|
||||
| Outcome | Reward |
|
||||
|---|---:|
|
||||
| all tests passed | 1.0 |
|
||||
| wrong answer | partial credit from test pass rate |
|
||||
| runtime error | -0.4 |
|
||||
| timeout | -0.6 |
|
||||
| syntax error | -0.8 |
|
||||
| missing code or policy violation | -1.0 |
|
||||
|
||||
Keep hidden tests separate from the prompts and report reward distributions by task difficulty. Otherwise the system can overfit public tests or produce groups with all-identical rewards.
|
||||
|
||||
## Safety boundary
|
||||
|
||||
The local verifier uses AST checks, a temporary working directory, an isolated Python process, a timeout, and POSIX resource limits. These are defense-in-depth measures, not a secure sandbox. Run hostile or externally submitted code inside a dedicated container or micro-VM with network disabled, a read-only filesystem, strict CPU/memory/process limits, and disposable credentials.
|
||||
|
||||
## Experimental protocol
|
||||
|
||||
Use the same held-out tasks and sampling parameters for three comparisons:
|
||||
|
||||
1. single-pass generation baseline;
|
||||
2. inference-time execution feedback and repair;
|
||||
3. execution-reward GRPO/CISPO followed by the same inference procedure.
|
||||
|
||||
Record pass@1, pass@k, test pass rate, compile rate, timeout rate, average attempts, average generation length, reward mean/std, KL to the reference model, and degenerate-group rate. Publish model results only after this protocol has been run and the raw logs and configuration have been retained for reproduction.
|
||||
|
||||
## Attribution
|
||||
|
||||
This repository is a derivative of [jingyaogong/minimind](https://github.com/jingyaogong/minimind) under Apache-2.0. The extension preserves upstream model names and checkpoint compatibility; presentation-layer identity prompts are kept neutral.
|
||||
|
|
@ -1,5 +1,6 @@
|
|||
datasets==3.6.0
|
||||
datasketch==1.6.4
|
||||
fastapi==0.115.6
|
||||
Flask==3.0.3
|
||||
Flask_Cors==4.0.0
|
||||
jieba==0.42.1
|
||||
|
|
@ -23,10 +24,11 @@ jinja2==3.1.2
|
|||
jsonlines==4.0.0
|
||||
trl==0.13.0
|
||||
ujson==5.1.0
|
||||
uvicorn==0.34.0
|
||||
wandb==0.18.3
|
||||
streamlit==1.50.0
|
||||
einops==0.8.1
|
||||
swanlab==0.7.11
|
||||
modelscope==1.37.0
|
||||
# torch==2.6.0
|
||||
# torchvision==0.21.0
|
||||
# torchvision==0.21.0
|
||||
|
|
|
|||
|
|
@ -0,0 +1,66 @@
|
|||
"""Evaluate JSONL code-generation candidates with deterministic tests."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
|
||||
from code_agent import CodeTask, aggregate_reports, evaluate_candidates
|
||||
|
||||
|
||||
def read_jsonl(path: str):
|
||||
with open(path, "r", encoding="utf-8") as handle:
|
||||
for line_number, line in enumerate(handle, start=1):
|
||||
if line.strip():
|
||||
try:
|
||||
yield json.loads(line)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ValueError(f"Invalid JSON at {path}:{line_number}: {exc}") from exc
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Execution-based evaluation for code generation")
|
||||
parser.add_argument("--tasks", required=True, help="JSONL file containing CodeTask records")
|
||||
parser.add_argument("--predictions", required=True, help="JSONL records with task_id and candidates")
|
||||
parser.add_argument("--output", help="Optional JSON report path")
|
||||
parser.add_argument("--ks", type=int, nargs="+", default=[1, 5, 10], help="pass@k values")
|
||||
parser.add_argument("--timeout", type=float, default=2.0, help="Timeout per candidate in seconds")
|
||||
parser.add_argument("--memory-mb", type=int, default=256, help="POSIX worker memory limit")
|
||||
args = parser.parse_args()
|
||||
|
||||
tasks = {record["task_id"]: CodeTask.from_dict(record) for record in read_jsonl(args.tasks)}
|
||||
predictions = {record["task_id"]: record["candidates"] for record in read_jsonl(args.predictions)}
|
||||
missing = sorted(set(tasks) - set(predictions))
|
||||
unknown = sorted(set(predictions) - set(tasks))
|
||||
if missing or unknown:
|
||||
raise ValueError(f"Task/prediction mismatch: missing={missing}, unknown={unknown}")
|
||||
|
||||
reports = [
|
||||
evaluate_candidates(
|
||||
task,
|
||||
predictions[task_id],
|
||||
ks=args.ks,
|
||||
timeout_seconds=args.timeout,
|
||||
memory_mb=args.memory_mb,
|
||||
)
|
||||
for task_id, task in tasks.items()
|
||||
]
|
||||
output = {
|
||||
"summary": aggregate_reports(reports),
|
||||
"tasks": [report.to_dict() for report in reports],
|
||||
"safety_note": (
|
||||
"The local verifier is for controlled experiments, not hostile code. "
|
||||
"Use a dedicated container or micro-VM for untrusted submissions."
|
||||
),
|
||||
}
|
||||
rendered = json.dumps(output, ensure_ascii=False, indent=2)
|
||||
if args.output:
|
||||
Path(args.output).write_text(rendered + "\n", encoding="utf-8")
|
||||
print(rendered)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,108 @@
|
|||
"""Generate reproducible multi-sample candidates for execution-based evaluation."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
|
||||
from code_agent import CodeTask, OpenAICompatibleGenerator
|
||||
|
||||
|
||||
def read_tasks(path):
|
||||
tasks = []
|
||||
with Path(path).open("r", encoding="utf-8") as handle:
|
||||
for line_number, line in enumerate(handle, start=1):
|
||||
if not line.strip():
|
||||
continue
|
||||
try:
|
||||
tasks.append(CodeTask.from_dict(json.loads(line)))
|
||||
except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
|
||||
raise ValueError(f"Invalid task at {path}:{line_number}: {exc}") from exc
|
||||
return tasks
|
||||
|
||||
|
||||
def build_prompt(task):
|
||||
return (
|
||||
"You are solving a Python algorithm task. Return exactly one ```python``` code block.\n"
|
||||
f"Required entry point: {task.entry_point}\n"
|
||||
f"Task:\n{task.prompt}\n"
|
||||
"Do not read files, access the network, spawn processes, or call unsafe builtins."
|
||||
)
|
||||
|
||||
|
||||
def generate_predictions(tasks, generator, samples_per_task):
|
||||
if samples_per_task < 1:
|
||||
raise ValueError("samples_per_task must be at least 1")
|
||||
return [
|
||||
{
|
||||
"task_id": task.task_id,
|
||||
"candidates": [generator(build_prompt(task)) for _ in range(samples_per_task)],
|
||||
}
|
||||
for task in tasks
|
||||
]
|
||||
|
||||
|
||||
def read_completed_task_ids(path):
|
||||
completed = set()
|
||||
if not Path(path).exists():
|
||||
return completed
|
||||
with Path(path).open("r", encoding="utf-8") as handle:
|
||||
for line_number, line in enumerate(handle, start=1):
|
||||
if not line.strip():
|
||||
continue
|
||||
try:
|
||||
record = json.loads(line)
|
||||
completed.add(str(record["task_id"]))
|
||||
except (KeyError, json.JSONDecodeError) as exc:
|
||||
raise ValueError(f"Invalid existing prediction at {path}:{line_number}: {exc}") from exc
|
||||
return completed
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Generate code candidates through an OpenAI-compatible API")
|
||||
parser.add_argument("--tasks", required=True, help="CodeTask JSONL file")
|
||||
parser.add_argument("--output", required=True, help="Prediction JSONL output path")
|
||||
parser.add_argument("--samples-per-task", type=int, default=10)
|
||||
parser.add_argument("--base-url", default="http://localhost:8998/v1")
|
||||
parser.add_argument("--api-key", default="not-needed")
|
||||
parser.add_argument("--model", default="minimind")
|
||||
parser.add_argument("--temperature", type=float, default=0.8)
|
||||
parser.add_argument("--max-tokens", type=int, default=512)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--open-thinking", action="store_true")
|
||||
parser.add_argument("--resume", action="store_true", help="Append missing tasks to an existing output file")
|
||||
args = parser.parse_args()
|
||||
|
||||
tasks = read_tasks(args.tasks)
|
||||
output_path = Path(args.output)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
completed = read_completed_task_ids(output_path) if args.resume else set()
|
||||
mode = "a" if args.resume else "w"
|
||||
generated_count = 0
|
||||
with output_path.open(mode, encoding="utf-8") as handle:
|
||||
for task_index, task in enumerate(tasks):
|
||||
if task.task_id in completed:
|
||||
print(f"[{task_index + 1}/{len(tasks)}] {task.task_id}: already complete", flush=True)
|
||||
continue
|
||||
generator = OpenAICompatibleGenerator(
|
||||
base_url=args.base_url,
|
||||
api_key=args.api_key,
|
||||
model=args.model,
|
||||
temperature=args.temperature,
|
||||
max_tokens=args.max_tokens,
|
||||
open_thinking=args.open_thinking,
|
||||
seed=args.seed + task_index * args.samples_per_task,
|
||||
)
|
||||
record = generate_predictions([task], generator, args.samples_per_task)[0]
|
||||
handle.write(json.dumps(record, ensure_ascii=False) + "\n")
|
||||
handle.flush()
|
||||
generated_count += args.samples_per_task
|
||||
print(f"[{task_index + 1}/{len(tasks)}] {task.task_id}: wrote {args.samples_per_task} candidates", flush=True)
|
||||
print(f"Generated {generated_count} new candidates -> {output_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,149 @@
|
|||
"""Convert compatible MBPP-sanitized assertions into MiniMind CodeTask splits."""
|
||||
|
||||
import argparse
|
||||
import ast
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
|
||||
from code_agent import CodeTask, VerificationStatus, verify_code
|
||||
|
||||
|
||||
def _literal(node):
|
||||
value = ast.literal_eval(node)
|
||||
json.dumps(value, ensure_ascii=False)
|
||||
return value
|
||||
|
||||
|
||||
def parse_assertion(source):
|
||||
statement = ast.parse(source.strip()).body
|
||||
if len(statement) != 1 or not isinstance(statement[0], ast.Assert):
|
||||
raise ValueError("test is not a single assert statement")
|
||||
expression = statement[0].test
|
||||
|
||||
if isinstance(expression, ast.Call):
|
||||
call, expected = expression, True
|
||||
elif isinstance(expression, ast.UnaryOp) and isinstance(expression.op, ast.Not) and isinstance(expression.operand, ast.Call):
|
||||
call, expected = expression.operand, False
|
||||
elif (
|
||||
isinstance(expression, ast.Compare)
|
||||
and len(expression.ops) == 1
|
||||
and isinstance(expression.ops[0], ast.Eq)
|
||||
and len(expression.comparators) == 1
|
||||
):
|
||||
left, right = expression.left, expression.comparators[0]
|
||||
if isinstance(left, ast.Call):
|
||||
call, expected = left, _literal(right)
|
||||
elif isinstance(right, ast.Call):
|
||||
call, expected = right, _literal(left)
|
||||
else:
|
||||
raise ValueError("equality does not compare a function call")
|
||||
else:
|
||||
raise ValueError("unsupported assertion shape")
|
||||
|
||||
if not isinstance(call.func, ast.Name):
|
||||
raise ValueError("entry point is not a direct function call")
|
||||
args = [_literal(argument) for argument in call.args]
|
||||
kwargs = {}
|
||||
for keyword in call.keywords:
|
||||
if keyword.arg is None:
|
||||
raise ValueError("expanded keyword arguments are unsupported")
|
||||
kwargs[keyword.arg] = _literal(keyword.value)
|
||||
return call.func.id, {"args": args, "kwargs": kwargs, "expected": expected}
|
||||
|
||||
|
||||
def convert_record(record):
|
||||
parsed = [parse_assertion(test) for test in record["test_list"]]
|
||||
entry_points = {entry_point for entry_point, _ in parsed}
|
||||
if len(entry_points) != 1:
|
||||
raise ValueError("tests use multiple entry points")
|
||||
payload = {
|
||||
"task_id": f"mbpp-{record['task_id']}",
|
||||
"prompt": record["prompt"].strip(),
|
||||
"entry_point": entry_points.pop(),
|
||||
"tests": [test for _, test in parsed],
|
||||
}
|
||||
task = CodeTask.from_dict(payload)
|
||||
verification = verify_code(record["code"], task)
|
||||
if verification.status != VerificationStatus.PASSED:
|
||||
raise ValueError(f"reference solution failed verifier: {verification.status.value}")
|
||||
return task.to_dict()
|
||||
|
||||
|
||||
def build_sft_record(record, task):
|
||||
return {
|
||||
"conversations": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": (
|
||||
"Solve the Python algorithm task below. Return one complete ```python``` code block.\n"
|
||||
f"Required entry point: {task['entry_point']}\n"
|
||||
f"Task:\n{task['prompt']}"
|
||||
),
|
||||
},
|
||||
{"role": "assistant", "content": f"```python\n{record['code'].strip()}\n```"},
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
def convert_split(records):
|
||||
converted = []
|
||||
sft_records = []
|
||||
skipped = Counter()
|
||||
for record in records:
|
||||
try:
|
||||
task = convert_record(record)
|
||||
converted.append(task)
|
||||
sft_records.append(build_sft_record(record, task))
|
||||
except (KeyError, SyntaxError, TypeError, ValueError) as exc:
|
||||
skipped[str(exc)] += 1
|
||||
return converted, sft_records, skipped
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Prepare verifier-compatible MBPP sanitized splits")
|
||||
parser.add_argument("--dataset", default="google-research-datasets/mbpp")
|
||||
parser.add_argument("--config", default="sanitized")
|
||||
parser.add_argument("--splits", nargs="+", default=["train", "validation", "test"])
|
||||
parser.add_argument("--output-dir", default="out/mbpp_sanitized")
|
||||
args = parser.parse_args()
|
||||
|
||||
from datasets import load_dataset
|
||||
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
summary = {}
|
||||
for split in args.splits:
|
||||
records = load_dataset(args.dataset, args.config, split=split)
|
||||
converted, sft_records, skipped = convert_split(records)
|
||||
output_path = output_dir / f"{split}.jsonl"
|
||||
with output_path.open("w", encoding="utf-8") as handle:
|
||||
for record in converted:
|
||||
handle.write(json.dumps(record, ensure_ascii=False) + "\n")
|
||||
sft_output = None
|
||||
if split == "train":
|
||||
sft_output = output_dir / "train_sft.jsonl"
|
||||
with sft_output.open("w", encoding="utf-8") as handle:
|
||||
for record in sft_records:
|
||||
handle.write(json.dumps(record, ensure_ascii=False) + "\n")
|
||||
summary[split] = {
|
||||
"source_count": len(records),
|
||||
"converted_count": len(converted),
|
||||
"skipped_count": len(records) - len(converted),
|
||||
"top_skip_reasons": skipped.most_common(10),
|
||||
"output": str(output_path),
|
||||
"sft_output": str(sft_output) if sft_output else None,
|
||||
}
|
||||
print(f"{split}: converted {len(converted)}/{len(records)} -> {output_path}")
|
||||
(output_dir / "summary.json").write_text(
|
||||
json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,146 @@
|
|||
"""Run the execution-feedback repair agent against an OpenAI-compatible model."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
|
||||
from code_agent import CodeTask, ExecutionFeedbackAgent, OpenAICompatibleGenerator
|
||||
|
||||
|
||||
TASK_SEED_STRIDE = 1000
|
||||
|
||||
|
||||
def read_tasks(path):
|
||||
tasks = []
|
||||
with Path(path).open("r", encoding="utf-8") as handle:
|
||||
for line_number, line in enumerate(handle, start=1):
|
||||
if not line.strip():
|
||||
continue
|
||||
try:
|
||||
tasks.append(CodeTask.from_dict(json.loads(line)))
|
||||
except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
|
||||
raise ValueError(f"Invalid task at {path}:{line_number}: {exc}") from exc
|
||||
return tasks
|
||||
|
||||
|
||||
def read_existing_results(path):
|
||||
output_path = Path(path)
|
||||
if not output_path.exists():
|
||||
return []
|
||||
payload = json.loads(output_path.read_text(encoding="utf-8"))
|
||||
results = payload.get("results")
|
||||
if not isinstance(results, list):
|
||||
raise ValueError(f"Invalid agent output at {path}: missing results list")
|
||||
seen = set()
|
||||
for result in results:
|
||||
task_id = str(result.get("task_id", ""))
|
||||
if not task_id or task_id in seen:
|
||||
raise ValueError(f"Invalid agent output at {path}: duplicate or missing task_id")
|
||||
seen.add(task_id)
|
||||
return results
|
||||
|
||||
|
||||
def build_output(results):
|
||||
return {
|
||||
"task_count": len(results),
|
||||
"success_count": sum(result["success"] for result in results),
|
||||
"results": results,
|
||||
}
|
||||
|
||||
|
||||
def write_output(path, output):
|
||||
output_path = Path(path)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
temporary_path = output_path.with_suffix(output_path.suffix + ".tmp")
|
||||
temporary_path.write_text(
|
||||
json.dumps(output, ensure_ascii=False, indent=2) + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
temporary_path.replace(output_path)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Generate, execute, and repair Python algorithm solutions")
|
||||
parser.add_argument("--tasks", required=True, help="CodeTask JSONL file")
|
||||
parser.add_argument("--task-id", help="Run only one task id")
|
||||
parser.add_argument("--output", help="Optional JSON output path")
|
||||
parser.add_argument("--base-url", default="http://localhost:8998/v1")
|
||||
parser.add_argument("--api-key", default="not-needed")
|
||||
parser.add_argument("--model", default="minimind")
|
||||
parser.add_argument("--temperature", type=float, default=0.7)
|
||||
parser.add_argument("--max-tokens", type=int, default=1024)
|
||||
parser.add_argument("--max-attempts", type=int, default=3)
|
||||
parser.add_argument("--timeout", type=float, default=2.0)
|
||||
parser.add_argument("--memory-mb", type=int, default=256)
|
||||
parser.add_argument("--seed", type=int, default=42, help="Base seed for reproducible task-level generation")
|
||||
parser.add_argument("--resume", action="store_true", help="Resume completed tasks from --output")
|
||||
parser.add_argument("--open-thinking", action="store_true")
|
||||
parser.add_argument("--reveal-test-details", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
indexed_tasks = list(enumerate(read_tasks(args.tasks)))
|
||||
if args.task_id:
|
||||
indexed_tasks = [(index, task) for index, task in indexed_tasks if task.task_id == args.task_id]
|
||||
if not indexed_tasks:
|
||||
raise ValueError(f"Unknown task id: {args.task_id}")
|
||||
|
||||
if args.resume and not args.output:
|
||||
parser.error("--resume requires --output")
|
||||
expected_task_ids = {task.task_id for _, task in indexed_tasks}
|
||||
existing_results = read_existing_results(args.output) if args.resume else []
|
||||
unexpected_task_ids = {result["task_id"] for result in existing_results} - expected_task_ids
|
||||
if unexpected_task_ids:
|
||||
raise ValueError(
|
||||
f"Existing output contains tasks not selected for this run: {sorted(unexpected_task_ids)}"
|
||||
)
|
||||
results_by_id = {result["task_id"]: result for result in existing_results}
|
||||
|
||||
for position, (task_index, task) in enumerate(indexed_tasks, start=1):
|
||||
if task.task_id in results_by_id:
|
||||
print(f"[{position}/{len(indexed_tasks)}] {task.task_id}: already complete", flush=True)
|
||||
continue
|
||||
generator = OpenAICompatibleGenerator(
|
||||
base_url=args.base_url,
|
||||
api_key=args.api_key,
|
||||
model=args.model,
|
||||
temperature=args.temperature,
|
||||
max_tokens=args.max_tokens,
|
||||
open_thinking=args.open_thinking,
|
||||
seed=args.seed + task_index * TASK_SEED_STRIDE,
|
||||
)
|
||||
agent = ExecutionFeedbackAgent(
|
||||
generator,
|
||||
max_attempts=args.max_attempts,
|
||||
timeout_seconds=args.timeout,
|
||||
memory_mb=args.memory_mb,
|
||||
reveal_test_details=args.reveal_test_details,
|
||||
)
|
||||
result = agent.run(task).to_dict()
|
||||
results_by_id[task.task_id] = result
|
||||
ordered_results = [
|
||||
results_by_id[selected_task.task_id]
|
||||
for _, selected_task in indexed_tasks
|
||||
if selected_task.task_id in results_by_id
|
||||
]
|
||||
if args.output:
|
||||
write_output(args.output, build_output(ordered_results))
|
||||
print(
|
||||
f"[{position}/{len(indexed_tasks)}] {task.task_id}: "
|
||||
f"success={result['success']} attempts={result['attempt_count']}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
results = [results_by_id[task.task_id] for _, task in indexed_tasks]
|
||||
output = build_output(results)
|
||||
rendered = json.dumps(output, ensure_ascii=False, indent=2)
|
||||
if args.output:
|
||||
write_output(args.output, output)
|
||||
print(rendered)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -56,6 +56,7 @@ class ChatRequest(BaseModel):
|
|||
stream: bool = True
|
||||
tools: list = Field(default_factory=list)
|
||||
open_thinking: bool = False
|
||||
seed: int | None = None
|
||||
chat_template_kwargs: dict = None
|
||||
|
||||
def get_open_thinking(self) -> bool:
|
||||
|
|
@ -102,7 +103,7 @@ def parse_response(text):
|
|||
return text.strip(), reasoning_content, tool_calls or None
|
||||
|
||||
|
||||
def generate_stream_response(messages, temperature, top_p, max_tokens, tools=None, open_thinking=False):
|
||||
def generate_stream_response(messages, temperature, top_p, max_tokens, tools=None, open_thinking=False, seed=None):
|
||||
try:
|
||||
new_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, tools=tools or None, open_thinking=open_thinking)
|
||||
inputs = tokenizer(new_prompt, return_tensors="pt", truncation=True).to(device)
|
||||
|
|
@ -112,6 +113,7 @@ def generate_stream_response(messages, temperature, top_p, max_tokens, tools=Non
|
|||
|
||||
def _generate():
|
||||
try:
|
||||
generator = torch.Generator(device=inputs.input_ids.device).manual_seed(seed) if seed is not None else None
|
||||
model.generate(
|
||||
inputs.input_ids,
|
||||
max_new_tokens=max_tokens,
|
||||
|
|
@ -121,7 +123,8 @@ def generate_stream_response(messages, temperature, top_p, max_tokens, tools=Non
|
|||
attention_mask=inputs.attention_mask,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
eos_token_id=tokenizer.eos_token_id,
|
||||
streamer=streamer
|
||||
streamer=streamer,
|
||||
generator=generator
|
||||
)
|
||||
except Exception as e:
|
||||
queue.put({"error": str(e)})
|
||||
|
|
@ -186,7 +189,8 @@ async def chat_completions(request: ChatRequest):
|
|||
top_p=request.top_p,
|
||||
max_tokens=request.max_tokens,
|
||||
tools=request.tools,
|
||||
open_thinking=request.get_open_thinking()
|
||||
open_thinking=request.get_open_thinking(),
|
||||
seed=request.seed
|
||||
)),
|
||||
media_type="text/event-stream"
|
||||
)
|
||||
|
|
@ -200,6 +204,7 @@ async def chat_completions(request: ChatRequest):
|
|||
)
|
||||
inputs = tokenizer(new_prompt, return_tensors="pt", truncation=True).to(device)
|
||||
with torch.no_grad():
|
||||
generator = torch.Generator(device=inputs["input_ids"].device).manual_seed(request.seed) if request.seed is not None else None
|
||||
generated_ids = model.generate(
|
||||
inputs["input_ids"],
|
||||
max_length=inputs["input_ids"].shape[1] + request.max_tokens,
|
||||
|
|
@ -208,7 +213,8 @@ async def chat_completions(request: ChatRequest):
|
|||
pad_token_id=tokenizer.pad_token_id,
|
||||
eos_token_id=tokenizer.eos_token_id,
|
||||
top_p=request.top_p,
|
||||
temperature=request.temperature
|
||||
temperature=request.temperature,
|
||||
generator=generator
|
||||
)
|
||||
answer = tokenizer.decode(generated_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
|
||||
content, reasoning_content, tool_calls = parse_response(answer)
|
||||
|
|
|
|||
|
|
@ -250,7 +250,7 @@ if not MODEL_PATHS:
|
|||
# 模型选择
|
||||
selected_model = st.sidebar.selectbox('Model', list(MODEL_PATHS.keys()), index=0)
|
||||
model_path = MODEL_PATHS[selected_model][0]
|
||||
slogan = f"我是 {MODEL_PATHS[selected_model][1]}，有什么可以帮你的？" if st.session_state.get('lang', 'en') == 'zh' else f"I am {MODEL_PATHS[selected_model][1]}, how can I help you?"
|
||||
slogan = "轻量级语言模型实验助手" if st.session_state.get('lang', 'en') == 'zh' else "Lightweight Language Model Assistant"
|
||||
|
||||
st.sidebar.markdown('<hr style="margin: 12px 0 16px 0;">', unsafe_allow_html=True)
|
||||
|
||||
|
|
@ -285,12 +285,9 @@ with st.sidebar.expander(get_text('tools')):
|
|||
if checked and len(st.session_state.selected_tools) < 4:
|
||||
st.session_state.selected_tools.append(name)
|
||||
|
||||
image_url = "https://www.modelscope.cn/api/v1/studio/gongjy/MiniMind/repo?Revision=master&FilePath=images%2Flogo2.png&View=true"
|
||||
|
||||
st.markdown(
|
||||
f'<div style="display: flex; flex-direction: column; align-items: center; text-align: center; margin: 0; padding: 0;">'
|
||||
'<div style="font-style: italic; font-weight: 900; margin: 0; padding-top: 4px; display: flex; align-items: center; justify-content: center; flex-wrap: wrap; width: 100%;">'
|
||||
f'<img src="{image_url}" style="width: 40px; height: 40px; "> '
|
||||
f'<span style="font-size: 26px; margin-left: 10px;">{slogan}</span>'
|
||||
'</div>'
|
||||
f'<span style="color: #bbb; font-style: italic; margin-top: 6px; margin-bottom: 10px;">{get_text("disclaimer")}</span>'
|
||||
|
|
@ -348,7 +345,7 @@ def main():
|
|||
setup_seed(random_seed)
|
||||
|
||||
tools = [t for t in TOOLS if t['function']['name'] in st.session_state.get('selected_tools', [])] or None
|
||||
sys_prompt = [] if tools else [{"role": "system", "content": "你是MiniMind，一个乐于助人、知识渊博的AI助手。请用完整且友好的方式回答用户问题。"}]
|
||||
sys_prompt = [] if tools else [{"role": "system", "content": "你是一个用于轻量语言模型实验的AI助手。请准确、完整且友好地回答用户问题，不要虚构开发者或个人身份。"}]
|
||||
st.session_state.chat_messages = sys_prompt + st.session_state.chat_messages[-(st.session_state.history_chat_num + 1):]
|
||||
template_kwargs = {"tokenize": False, "add_generation_prompt": True}
|
||||
if st.session_state.get('enable_thinking', False):
|
||||
|
|
|
|||
|
|
@ -0,0 +1 @@
|
|||
{"task_id":"add-two","candidates":["def add(a, b):\n return a + b","def add(a, b):\n return a - b"]}
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
{"task_id":"reverse-string","candidates":["def reverse_string(text):\n return text[::-1]","def reverse_string(text):\n return text"]}
|
||||
{"task_id":"valid-parentheses","candidates":["def valid_parentheses(text):\n pairs = {')': '(', ']': '[', '}': '{'}\n stack = []\n for char in text:\n if char in '([{':\n stack.append(char)\n elif not stack or stack.pop() != pairs[char]:\n return False\n return not stack","def valid_parentheses(text):\n return True"]}
|
||||
{"task_id":"binary-search","candidates":["def binary_search(values, target):\n left, right = 0, len(values) - 1\n while left <= right:\n mid = (left + right) // 2\n if values[mid] == target:\n return mid\n if values[mid] < target:\n left = mid + 1\n else:\n right = mid - 1\n return -1","def binary_search(values, target):\n return -1"]}
|
||||
{"task_id":"max-subarray","candidates":["def max_subarray(values):\n best = current = values[0]\n for value in values[1:]:\n current = max(value, current + value)\n best = max(best, current)\n return best","def max_subarray(values):\n return sum(values)"]}
|
||||
|
|
@ -0,0 +1 @@
|
|||
{"task_id":"add-two","prompt":"Implement add(a, b) and return the sum.","entry_point":"add","tests":[{"args":[1,2],"expected":3},{"args":[-5,5],"expected":0}]}
|
||||
|
|
@ -0,0 +1,47 @@
|
|||
from code_agent import CodeTask, ExecutionFeedbackAgent, TestCase as CodeTestCase
|
||||
|
||||
|
||||
TASK = CodeTask(
|
||||
task_id="maximum",
|
||||
prompt="Implement maximum(values), returning the largest integer.",
|
||||
entry_point="maximum",
|
||||
tests=(
|
||||
CodeTestCase(args=([1, 9, 3],), expected=9),
|
||||
CodeTestCase(args=([-5, -2, -8],), expected=-2),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class ScriptedGenerator:
|
||||
def __init__(self, responses):
|
||||
self.responses = iter(responses)
|
||||
self.prompts = []
|
||||
|
||||
def __call__(self, prompt):
|
||||
self.prompts.append(prompt)
|
||||
return next(self.responses)
|
||||
|
||||
|
||||
def test_agent_repairs_solution_from_hidden_test_feedback():
|
||||
generator = ScriptedGenerator(
|
||||
[
|
||||
"```python\ndef maximum(values):\n return values[0]\n```",
|
||||
"```python\ndef maximum(values):\n return max(values)\n```",
|
||||
]
|
||||
)
|
||||
result = ExecutionFeedbackAgent(generator, max_attempts=3).run(TASK)
|
||||
|
||||
assert result.success
|
||||
assert len(result.attempts) == 2
|
||||
assert result.attempts[0].verification.pass_rate == 0.0
|
||||
assert "hidden test failed" in generator.prompts[1]
|
||||
assert "expected 9" not in generator.prompts[1]
|
||||
assert result.final.reward.total == 1.0
|
||||
|
||||
|
||||
def test_agent_stops_at_attempt_limit():
|
||||
generator = ScriptedGenerator(["def maximum(values):\n return 0"] * 2)
|
||||
result = ExecutionFeedbackAgent(generator, max_attempts=2).run(TASK)
|
||||
|
||||
assert not result.success
|
||||
assert len(result.attempts) == 2
|
||||
|
|
@ -0,0 +1,47 @@
|
|||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from code_agent import OpenAICompatibleGenerator
|
||||
|
||||
|
||||
class FakeCompletions:
|
||||
def __init__(self, content):
|
||||
self.content = content
|
||||
self.kwargs = None
|
||||
|
||||
def create(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
message = SimpleNamespace(content=self.content)
|
||||
return SimpleNamespace(choices=[SimpleNamespace(message=message)])
|
||||
|
||||
|
||||
def test_openai_backend_builds_non_streaming_request():
|
||||
completions = FakeCompletions("```python\ndef solve():\n return 1\n```")
|
||||
client = SimpleNamespace(chat=SimpleNamespace(completions=completions))
|
||||
generator = OpenAICompatibleGenerator(
|
||||
client=client,
|
||||
model="local-minimind",
|
||||
temperature=0.2,
|
||||
max_tokens=128,
|
||||
open_thinking=True,
|
||||
seed=17,
|
||||
)
|
||||
|
||||
assert generator("solve it").startswith("```python")
|
||||
assert completions.kwargs["messages"][0]["content"] == "solve it"
|
||||
assert completions.kwargs["stream"] is False
|
||||
assert completions.kwargs["extra_body"] == {"open_thinking": True}
|
||||
assert completions.kwargs["seed"] == 17
|
||||
|
||||
generator("solve another")
|
||||
assert completions.kwargs["seed"] == 18
|
||||
|
||||
|
||||
def test_openai_backend_rejects_empty_completion():
|
||||
completions = FakeCompletions("")
|
||||
client = SimpleNamespace(chat=SimpleNamespace(completions=completions))
|
||||
generator = OpenAICompatibleGenerator(client=client)
|
||||
|
||||
with pytest.raises(RuntimeError, match="empty completion"):
|
||||
generator("solve it")
|
||||
|
|
@ -0,0 +1,54 @@
|
|||
import pytest
|
||||
|
||||
from code_agent import (
|
||||
CodeTask,
|
||||
TestCase as CodeTestCase,
|
||||
VerificationStatus,
|
||||
aggregate_reports,
|
||||
estimate_pass_at_k,
|
||||
evaluate_candidates,
|
||||
)
|
||||
|
||||
|
||||
TASK = CodeTask(
|
||||
task_id="square",
|
||||
prompt="Implement square(x).",
|
||||
entry_point="square",
|
||||
tests=(
|
||||
CodeTestCase(args=(2,), expected=4),
|
||||
CodeTestCase(args=(-3,), expected=9),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_pass_at_k_estimator():
|
||||
assert estimate_pass_at_k(10, 0, 1) == 0.0
|
||||
assert estimate_pass_at_k(10, 1, 10) == 1.0
|
||||
assert estimate_pass_at_k(10, 2, 1) == pytest.approx(0.2)
|
||||
with pytest.raises(ValueError):
|
||||
estimate_pass_at_k(2, 1, 3)
|
||||
|
||||
|
||||
def test_evaluation_reports_execution_metrics():
|
||||
candidates = [
|
||||
"def square(x):\n return x * x",
|
||||
"def square(x):\n return x + x",
|
||||
"def square(x)\n return x * x",
|
||||
]
|
||||
report = evaluate_candidates(TASK, candidates, ks=(1, 3))
|
||||
|
||||
assert report.correct_count == 1
|
||||
assert report.compile_rate == pytest.approx(2 / 3)
|
||||
assert report.execution_rate == pytest.approx(2 / 3)
|
||||
assert report.test_pass_rate == pytest.approx(0.5)
|
||||
assert report.pass_at_k["pass@1"] == pytest.approx(1 / 3)
|
||||
assert report.pass_at_k["pass@3"] == 1.0
|
||||
assert report.results[2].status == VerificationStatus.SYNTAX_ERROR
|
||||
|
||||
summary = aggregate_reports([report])
|
||||
assert summary["candidate_count"] == 3
|
||||
assert summary["status_counts"] == {
|
||||
"passed": 1,
|
||||
"syntax_error": 1,
|
||||
"wrong_answer": 1,
|
||||
}
|
||||
|
|
@ -0,0 +1,25 @@
|
|||
import pytest
|
||||
|
||||
from code_agent import CodeTask, TestCase as CodeTestCase, score_code_batch
|
||||
|
||||
|
||||
TASK = CodeTask(
|
||||
task_id="identity",
|
||||
prompt="Implement identity(x).",
|
||||
entry_point="identity",
|
||||
tests=(CodeTestCase(args=(7,), expected=7),),
|
||||
)
|
||||
|
||||
|
||||
def test_scores_task_major_grouped_responses():
|
||||
rewards = score_code_batch(
|
||||
[TASK],
|
||||
["def identity(x):\n return x", "def identity(x):\n return 0"],
|
||||
num_generations=2,
|
||||
)
|
||||
assert [reward.total for reward in rewards] == [1.0, 0.0]
|
||||
|
||||
|
||||
def test_rejects_misaligned_group_size():
|
||||
with pytest.raises(ValueError, match="expected 2 responses"):
|
||||
score_code_batch([TASK], ["def identity(x): return x"], num_generations=2)
|
||||
|
|
@ -0,0 +1,31 @@
|
|||
import json
|
||||
|
||||
from code_agent import CodeRLDataset, collate_code_rl
|
||||
|
||||
|
||||
class FakeTokenizer:
|
||||
def apply_chat_template(self, messages, **kwargs):
|
||||
return messages[0]["content"] + f"\nthinking={kwargs['open_thinking']}"
|
||||
|
||||
|
||||
def test_code_rl_dataset_builds_prompt_and_preserves_tests(tmp_path):
|
||||
path = tmp_path / "tasks.jsonl"
|
||||
path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"task_id": "inc",
|
||||
"prompt": "Implement inc(x).",
|
||||
"entry_point": "inc",
|
||||
"tests": [{"args": [1], "expected": 2}],
|
||||
}
|
||||
)
|
||||
+ "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
dataset = CodeRLDataset(path, FakeTokenizer(), thinking_ratio=0.0)
|
||||
sample = dataset[0]
|
||||
|
||||
assert "Required entry point: inc" in sample["prompt"]
|
||||
assert sample["task"].tests[0].expected == 2
|
||||
batch = collate_code_rl([sample])
|
||||
assert batch["tasks"][0].task_id == "inc"
|
||||
|
|
@ -0,0 +1,57 @@
|
|||
from code_agent import (
|
||||
CodeTask,
|
||||
TestCase as CodeTestCase,
|
||||
VerificationStatus,
|
||||
extract_python_code,
|
||||
score_verification,
|
||||
verify_code,
|
||||
)
|
||||
|
||||
|
||||
TASK = CodeTask(
|
||||
task_id="add-two",
|
||||
prompt="Implement add(a, b).",
|
||||
entry_point="add",
|
||||
tests=(
|
||||
CodeTestCase(args=(1, 2), expected=3),
|
||||
CodeTestCase(args=(5, 0), expected=5),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_extracts_fenced_code_after_thinking():
|
||||
response = "<think>reasoning</think>\n```python\ndef add(a, b):\n return a + b\n```"
|
||||
assert extract_python_code(response).startswith("def add")
|
||||
|
||||
|
||||
def test_verifier_accepts_correct_solution():
|
||||
result = verify_code("def add(a, b):\n return a + b", TASK)
|
||||
assert result.status == VerificationStatus.PASSED
|
||||
assert result.pass_rate == 1.0
|
||||
assert score_verification(result).total == 1.0
|
||||
|
||||
|
||||
def test_verifier_reports_partial_credit_and_feedback():
|
||||
result = verify_code("def add(a, b):\n return a - b", TASK)
|
||||
assert result.status == VerificationStatus.WRONG_ANSWER
|
||||
assert result.passed_tests == 1
|
||||
assert "expected 3, got -1" in result.feedback
|
||||
assert score_verification(result).total == 0.4
|
||||
|
||||
|
||||
def test_verifier_reports_syntax_error():
|
||||
result = verify_code("def add(a, b)\n return a + b", TASK)
|
||||
assert result.status == VerificationStatus.SYNTAX_ERROR
|
||||
assert score_verification(result).total == -0.8
|
||||
|
||||
|
||||
def test_verifier_blocks_unsafe_import():
|
||||
result = verify_code("import os\ndef add(a, b):\n return a + b", TASK)
|
||||
assert result.status == VerificationStatus.POLICY_VIOLATION
|
||||
assert "import 'os' is not allowed" in result.feedback
|
||||
|
||||
|
||||
def test_verifier_times_out_infinite_loop():
|
||||
result = verify_code("def add(a, b):\n while True:\n pass", TASK, timeout_seconds=0.2)
|
||||
assert result.status == VerificationStatus.TIMEOUT
|
||||
assert score_verification(result).total == -0.6
|
||||
|
|
@ -0,0 +1,37 @@
|
|||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from scripts.generate_codegen import build_prompt, generate_predictions, read_completed_task_ids
|
||||
|
||||
|
||||
def test_generate_predictions_groups_samples_by_task():
|
||||
task = SimpleNamespace(task_id="add", entry_point="add", prompt="Implement add(a, b).")
|
||||
prompts = []
|
||||
|
||||
def generate(prompt):
|
||||
prompts.append(prompt)
|
||||
return f"candidate-{len(prompts)}"
|
||||
|
||||
records = generate_predictions([task], generate, samples_per_task=3)
|
||||
|
||||
assert records == [{"task_id": "add", "candidates": ["candidate-1", "candidate-2", "candidate-3"]}]
|
||||
assert all("Required entry point: add" in prompt for prompt in prompts)
|
||||
assert "exactly one ```python``` code block" in build_prompt(task)
|
||||
|
||||
|
||||
def test_generate_predictions_rejects_empty_sample_count():
|
||||
with pytest.raises(ValueError, match="at least 1"):
|
||||
generate_predictions([], lambda prompt: prompt, samples_per_task=0)
|
||||
|
||||
|
||||
def test_read_completed_task_ids_supports_resume(tmp_path):
|
||||
output = tmp_path / "predictions.jsonl"
|
||||
output.write_text(
|
||||
'{"task_id":"one","candidates":["a"]}\n'
|
||||
'{"task_id":"two","candidates":["b"]}\n',
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
assert read_completed_task_ids(output) == {"one", "two"}
|
||||
assert read_completed_task_ids(tmp_path / "missing.jsonl") == set()
|
||||
|
|
@ -0,0 +1,50 @@
|
|||
import pytest
|
||||
|
||||
from scripts.prepare_mbpp import build_sft_record, convert_record, convert_split, parse_assertion
|
||||
|
||||
|
||||
def test_parse_mbpp_equality_and_boolean_assertions():
|
||||
assert parse_assertion("assert add(1, 2) == 3") == (
|
||||
"add",
|
||||
{"args": [1, 2], "kwargs": {}, "expected": 3},
|
||||
)
|
||||
assert parse_assertion("assert is_even(4)")[-1]["expected"] is True
|
||||
assert parse_assertion("assert not is_even(3)")[-1]["expected"] is False
|
||||
|
||||
|
||||
def test_convert_record_validates_reference_solution():
|
||||
record = {
|
||||
"task_id": 1,
|
||||
"prompt": "Write a function that adds two integers.",
|
||||
"code": "def add(a, b):\n return a + b",
|
||||
"test_list": ["assert add(1, 2) == 3", "assert add(-1, 1) == 0"],
|
||||
}
|
||||
|
||||
converted = convert_record(record)
|
||||
|
||||
assert converted["task_id"] == "mbpp-1"
|
||||
assert converted["entry_point"] == "add"
|
||||
assert len(converted["tests"]) == 2
|
||||
|
||||
sft = build_sft_record(record, converted)
|
||||
assert "Required entry point: add" in sft["conversations"][0]["content"]
|
||||
assert "def add(a, b)" in sft["conversations"][1]["content"]
|
||||
|
||||
|
||||
def test_convert_split_keeps_tasks_and_sft_records_aligned():
|
||||
record = {
|
||||
"task_id": 2,
|
||||
"prompt": "Return the input integer.",
|
||||
"code": "def identity(value):\n return value",
|
||||
"test_list": ["assert identity(7) == 7"],
|
||||
}
|
||||
|
||||
tasks, sft_records, skipped = convert_split([record])
|
||||
|
||||
assert len(tasks) == len(sft_records) == 1
|
||||
assert not skipped
|
||||
|
||||
|
||||
def test_parse_mbpp_rejects_nonliteral_expectation():
|
||||
with pytest.raises(ValueError):
|
||||
parse_assertion("assert add(1, 2) == expected")
|
||||
|
|
@ -0,0 +1,39 @@
|
|||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from scripts.run_code_agent import build_output, read_existing_results, write_output
|
||||
|
||||
|
||||
def test_agent_output_round_trip(tmp_path):
|
||||
output_path = tmp_path / "agent.json"
|
||||
results = [
|
||||
{"task_id": "task-a", "success": True, "attempt_count": 1, "attempts": []},
|
||||
{"task_id": "task-b", "success": False, "attempt_count": 3, "attempts": []},
|
||||
]
|
||||
|
||||
output = build_output(results)
|
||||
write_output(output_path, output)
|
||||
|
||||
assert json.loads(output_path.read_text(encoding="utf-8")) == output
|
||||
assert read_existing_results(output_path) == results
|
||||
assert output["task_count"] == 2
|
||||
assert output["success_count"] == 1
|
||||
|
||||
|
||||
def test_agent_resume_rejects_duplicate_task_ids(tmp_path):
|
||||
output_path = tmp_path / "agent.json"
|
||||
output_path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"results": [
|
||||
{"task_id": "task-a", "success": False},
|
||||
{"task_id": "task-a", "success": True},
|
||||
]
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="duplicate or missing task_id"):
|
||||
read_existing_results(output_path)
|
||||
|
|
@ -24,6 +24,7 @@ from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
|||
from dataset.lm_dataset import RLAIFDataset
|
||||
from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
|
||||
from trainer.rollout_engine import create_rollout_engine
|
||||
from code_agent import CodeRLDataset, collate_code_rl, score_code_batch
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
|
@ -34,7 +35,17 @@ def rep_penalty(text, n=3, cap=0.5):
|
|||
return min(cap, (len(grams) - len(set(grams))) * cap * 2 / len(grams)) if grams else 0.0
|
||||
|
||||
|
||||
def calculate_rewards(prompts, responses, reward_model):
|
||||
def calculate_rewards(prompts, responses, reward_model, tasks=None):
|
||||
if args.reward_type == "code":
|
||||
code_rewards = score_code_batch(
|
||||
tasks,
|
||||
responses,
|
||||
num_generations=args.num_generations,
|
||||
timeout_seconds=args.code_timeout,
|
||||
memory_mb=args.code_memory_mb,
|
||||
)
|
||||
return torch.tensor([reward.total for reward in code_rewards], device=args.device)
|
||||
|
||||
rewards = torch.zeros(len(responses), device=args.device)
|
||||
|
||||
with torch.no_grad():
|
||||
|
|
@ -92,7 +103,7 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
|
|||
full_mask = (outputs != tokenizer.pad_token_id).long()
|
||||
logp_pos = prompt_lens.unsqueeze(1) - 1 + torch.arange(completion_ids.size(1), device=args.device).unsqueeze(0)
|
||||
|
||||
rewards = calculate_rewards(prompts, completions, reward_model).to(args.device) # [B*num_gen]
|
||||
rewards = calculate_rewards(prompts, completions, reward_model, batch.get('tasks')).to(args.device) # [B*num_gen]
|
||||
|
||||
model_unwrapped = model.module if isinstance(model, DistributedDataParallel) else model
|
||||
with autocast_ctx:
|
||||
|
|
@ -119,6 +130,7 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
|
|||
Logger('='*100)
|
||||
|
||||
grouped_rewards = rewards.view(-1, args.num_generations) # [B, num_gen]
|
||||
degenerate_group_rate = (grouped_rewards.std(dim=1, unbiased=False) < 1e-6).float().mean()
|
||||
mean_r = grouped_rewards.mean(dim=1).repeat_interleave(args.num_generations) # [B*num_gen]
|
||||
std_r = grouped_rewards.std(dim=1, unbiased=False).repeat_interleave(args.num_generations) # [B*num_gen]
|
||||
advantages = (rewards - mean_r) / (std_r + 1e-4) # [B*num_gen]
|
||||
|
|
@ -156,13 +168,14 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
|
|||
current_aux_loss = aux_loss.item()
|
||||
avg_reward_val = rewards.mean().item()
|
||||
avg_len_val = completion_mask.sum(dim=1).float().mean().item()
|
||||
kl_ref_val = ((ref_per_token_logps - per_token_logps) * completion_mask).sum().item() / max(completion_mask.sum().item(), 1)
|
||||
kl_ref_val = (per_token_kl * completion_mask).sum().item() / max(completion_mask.sum().item(), 1)
|
||||
advantages_mean_val = advantages.mean().item()
|
||||
advantages_std_val = advantages.std().item()
|
||||
current_lr = optimizer.param_groups[0]['lr']
|
||||
|
||||
Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), '
|
||||
f'Reward: {avg_reward_val:.4f}, KL_ref: {kl_ref_val:.4f}, '
|
||||
f'Degenerate Groups: {degenerate_group_rate.item():.4f}, '
|
||||
f'Adv Std: {advantages_std_val:.4f}, Adv Mean: {advantages_mean_val:.4f}, '
|
||||
f'Actor Loss: {policy_loss_val:.4f}, Avg Response Len: {avg_len_val:.2f}, Learning Rate: {current_lr:.8f}')
|
||||
|
||||
|
|
@ -174,6 +187,7 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
|
|||
"advantages_mean": advantages_mean_val,
|
||||
"policy_loss": policy_loss_val,
|
||||
"avg_response_len": avg_len_val,
|
||||
"degenerate_group_rate": degenerate_group_rate.item(),
|
||||
"learning_rate": current_lr
|
||||
})
|
||||
|
||||
|
|
@ -230,6 +244,9 @@ if __name__ == "__main__":
|
|||
parser.add_argument("--epsilon_high", type=float, default=5.0, help="epsilon上界")
|
||||
parser.add_argument('--from_weight', default='full_sft', type=str, help="基于哪个权重训练")
|
||||
parser.add_argument("--reward_model_path", type=str, default="../../internlm2-1_8b-reward", help="Reward模型路径")
|
||||
parser.add_argument("--reward_type", type=str, default="reward_model", choices=["reward_model", "code"], help="Reward source")
|
||||
parser.add_argument("--code_timeout", type=float, default=2.0, help="Code execution timeout per candidate")
|
||||
parser.add_argument("--code_memory_mb", type=int, default=256, help="POSIX memory limit for each code worker")
|
||||
parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
|
||||
parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
|
||||
parser.add_argument("--wandb_project", type=str, default="MiniMind-GRPO", help="wandb项目名")
|
||||
|
|
@ -276,7 +293,10 @@ if __name__ == "__main__":
|
|||
ref_model, _ = init_model(lm_config, base_weight, device=args.device)
|
||||
ref_model = ref_model.eval().requires_grad_(False)
|
||||
# Reward模型
|
||||
reward_model = LMForRewardModel(args.reward_model_path, device=args.device, dtype=torch.float16)
|
||||
reward_model = (
|
||||
LMForRewardModel(args.reward_model_path, device=args.device, dtype=torch.float16)
|
||||
if args.reward_type == "reward_model" else None
|
||||
)
|
||||
# Rollout引擎（可插拔替换，只负责 policy 推理）
|
||||
rollout_engine = create_rollout_engine(
|
||||
engine_type=args.rollout_engine,
|
||||
|
|
@ -289,7 +309,11 @@ if __name__ == "__main__":
|
|||
sglang_shared_path=args.sglang_shared_path,
|
||||
)
|
||||
# 数据和优化器
|
||||
train_ds = RLAIFDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len, thinking_ratio=args.thinking_ratio)
|
||||
train_ds = (
|
||||
CodeRLDataset(args.data_path, tokenizer, thinking_ratio=args.thinking_ratio)
|
||||
if args.reward_type == "code"
|
||||
else RLAIFDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len, thinking_ratio=args.thinking_ratio)
|
||||
)
|
||||
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
|
||||
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
|
||||
loader_for_count = DataLoader(train_ds, batch_size=args.batch_size, sampler=train_sampler)
|
||||
|
|
@ -321,7 +345,8 @@ if __name__ == "__main__":
|
|||
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
|
||||
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
|
||||
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
|
||||
collate_fn = collate_code_rl if args.reward_type == "code" else None
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True, collate_fn=collate_fn)
|
||||
if skip > 0:
|
||||
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
|
||||
grpo_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, reward_model, start_step, wandb, use_sglang = (args.rollout_engine == "sglang"))
|
||||
|
|
@ -331,4 +356,4 @@ if __name__ == "__main__":
|
|||
# ========== 9. 清理分布进程 ==========
|
||||
if dist.is_initialized():
|
||||
dist.barrier()
|
||||
dist.destroy_process_group()
|
||||
dist.destroy_process_group()
|
||||
|
|
|
|||
Loading…
Reference in New Issue