|
|
"""
|
|
|
AIME数学题目生成器
|
|
|
|
|
|
使用HelloAgents框架生成AIME风格的数学题目
|
|
|
"""
|
|
|
|
|
|
import json
|
|
|
import os
|
|
|
import time
|
|
|
import random
|
|
|
from typing import List, Dict, Any, Optional
|
|
|
from datetime import datetime
|
|
|
from tqdm import tqdm
|
|
|
from hello_agents import SimpleAgent
|
|
|
from hello_agents import HelloAgentsLLM
|
|
|
from datasets import load_dataset
|
|
|
|
|
|
|
|
|
class AIMEGenerator:
|
|
|
"""AIME题目生成器"""
|
|
|
|
|
|
# AIME题目生成提示词（英文）
|
|
|
GENERATION_PROMPT = """You are a professional mathematics competition problem designer, skilled in creating AIME (American Invitational Mathematics Examination) style problems.
|
|
|
|
|
|
AIME Problem Characteristics:
|
|
|
1. Answer: An integer between 0 and 999
|
|
|
2. Topics: Algebra, Geometry, Number Theory, Combinatorics, Probability, etc.
|
|
|
3. Style: Requires multi-step reasoning, but no advanced theory
|
|
|
4. Difficulty: Medium to hard (similar to AIME problems 6-9)
|
|
|
|
|
|
Please generate an AIME-style mathematics problem, including:
|
|
|
1. Problem statement (clear and complete)
|
|
|
2. Answer (an integer between 0 and 999)
|
|
|
3. Detailed solution (including all reasoning steps)
|
|
|
4. Topic classification (Algebra/Geometry/Number Theory/Combinatorics/Probability)
|
|
|
|
|
|
Please output in the following JSON format, avoid using special escape characters in JSON:
|
|
|
```json
|
|
|
{
|
|
|
"problem": "Problem statement in English",
|
|
|
"answer": 123,
|
|
|
"solution": "Detailed solution steps in English",
|
|
|
"topic": "Algebra"
|
|
|
}
|
|
|
```
|
|
|
"""
|
|
|
|
|
|
def __init__(
|
|
|
self,
|
|
|
llm: HelloAgentsLLM = None,
|
|
|
delay_seconds: float = 1.0,
|
|
|
use_reference_examples: bool = True,
|
|
|
reference_dataset: str = "TianHongZXY/aime-1983-2025"
|
|
|
):
|
|
|
"""
|
|
|
初始化生成器
|
|
|
|
|
|
Args:
|
|
|
llm: LLM实例（可选）
|
|
|
delay_seconds: 每次生成之间的延迟（秒），避免API速率限制
|
|
|
use_reference_examples: 是否使用真题作为参考样例
|
|
|
reference_dataset: 参考数据集名称，默认使用TianHongZXY/aime-1983-2025（900+道题）
|
|
|
"""
|
|
|
# 如果没有提供llm，创建默认的HelloAgentsLLM
|
|
|
if llm is None:
|
|
|
self.llm = HelloAgentsLLM()
|
|
|
else:
|
|
|
self.llm = llm
|
|
|
|
|
|
self.agent = SimpleAgent(
|
|
|
name="AIME Generator",
|
|
|
llm=self.llm,
|
|
|
system_prompt="你是一位专业的数学竞赛题目设计专家。"
|
|
|
)
|
|
|
self.delay_seconds = delay_seconds
|
|
|
self.use_reference_examples = use_reference_examples
|
|
|
self.reference_examples = []
|
|
|
|
|
|
# 加载参考样例
|
|
|
if use_reference_examples:
|
|
|
try:
|
|
|
print(f"📚 加载AIME真题数据集: {reference_dataset}")
|
|
|
# 尝试不同的split
|
|
|
try:
|
|
|
dataset = load_dataset(reference_dataset, split="train")
|
|
|
except:
|
|
|
dataset = load_dataset(reference_dataset, split="test")
|
|
|
|
|
|
# 加载所有题目作为参考
|
|
|
self.reference_examples = list(dataset)
|
|
|
print(f" ✓ 已加载 {len(self.reference_examples)} 道参考题目")
|
|
|
|
|
|
# 统计年份分布（如果有year字段）
|
|
|
year_counts = {}
|
|
|
for item in self.reference_examples:
|
|
|
year = item.get('year')
|
|
|
if year:
|
|
|
year_counts[year] = year_counts.get(year, 0) + 1
|
|
|
|
|
|
if year_counts:
|
|
|
year_range = f"{min(year_counts.keys())}-{max(year_counts.keys())}"
|
|
|
print(f" ℹ️ 年份范围: {year_range}")
|
|
|
|
|
|
except Exception as e:
|
|
|
print(f" ⚠️ 加载参考样例失败: {e}")
|
|
|
print(f" ℹ️ 将使用默认提示词生成")
|
|
|
self.use_reference_examples = False
|
|
|
|
|
|
def generate_single(self, max_retries: int = 3) -> Dict[str, Any]:
|
|
|
"""
|
|
|
生成单个题目
|
|
|
|
|
|
Args:
|
|
|
max_retries: 最大重试次数
|
|
|
|
|
|
Returns:
|
|
|
题目数据
|
|
|
"""
|
|
|
# 构建提示词
|
|
|
prompt = self._build_prompt()
|
|
|
|
|
|
for attempt in range(max_retries):
|
|
|
try:
|
|
|
response = self.agent.run(prompt)
|
|
|
return self._parse_response(response)
|
|
|
except Exception as e:
|
|
|
if attempt < max_retries - 1:
|
|
|
tqdm.write(f"⚠️ 生成失败（尝试 {attempt + 1}/{max_retries}），{self.delay_seconds}秒后重试...")
|
|
|
time.sleep(self.delay_seconds)
|
|
|
else:
|
|
|
tqdm.write(f"❌ 生成失败，已达最大重试次数: {e}")
|
|
|
return self._get_default_problem()
|
|
|
|
|
|
def _build_prompt(self) -> str:
|
|
|
"""构建生成提示词"""
|
|
|
if not self.use_reference_examples or not self.reference_examples:
|
|
|
return self.GENERATION_PROMPT
|
|
|
|
|
|
# 随机选择一个参考样例
|
|
|
example = random.choice(self.reference_examples)
|
|
|
example_problem = example.get('problem', 'Example problem')
|
|
|
example_answer = example.get('answer', 0)
|
|
|
|
|
|
# 构建带参考样例的提示词（英文）
|
|
|
prompt = f"""You are a professional mathematics competition problem designer, skilled in creating AIME (American Invitational Mathematics Examination) style problems.
|
|
|
|
|
|
【Reference Example】(For style reference only, please generate a completely different problem)
|
|
|
Problem: {example_problem}
|
|
|
Answer: {example_answer}
|
|
|
|
|
|
AIME Problem Characteristics:
|
|
|
1. Answer: An integer between 0 and 999
|
|
|
2. Topics: Algebra, Geometry, Number Theory, Combinatorics, Probability, etc.
|
|
|
3. Style: Requires multi-step reasoning, but no advanced theory
|
|
|
4. Difficulty: Medium to hard (similar to AIME problems 6-9)
|
|
|
|
|
|
Please generate a **completely different** AIME-style mathematics problem, including:
|
|
|
1. Problem statement (clear and complete, different from the reference)
|
|
|
2. Answer (an integer between 0 and 999, different from the reference)
|
|
|
3. Detailed solution (including all reasoning steps)
|
|
|
4. Topic classification (Algebra/Geometry/Number Theory/Combinatorics/Probability)
|
|
|
|
|
|
Please output in the following JSON format, avoid using special escape characters in JSON:
|
|
|
```json
|
|
|
{{
|
|
|
"problem": "Problem statement in English",
|
|
|
"answer": 123,
|
|
|
"solution": "Detailed solution steps in English",
|
|
|
"topic": "Algebra"
|
|
|
}}
|
|
|
```
|
|
|
|
|
|
Important Notes:
|
|
|
- **Must generate a completely different problem from the reference**
|
|
|
- You can reference the style, but do not copy the content
|
|
|
- Ensure the problem is creative and original
|
|
|
"""
|
|
|
return prompt
|
|
|
|
|
|
def _parse_response(self, response: str) -> Dict[str, Any]:
|
|
|
"""解析LLM响应（支持LaTeX数学公式）"""
|
|
|
import re
|
|
|
|
|
|
# 提取JSON部分
|
|
|
if "```json" in response:
|
|
|
json_str = response.split("```json")[1].split("```")[0].strip()
|
|
|
elif "```" in response:
|
|
|
json_str = response.split("```")[1].split("```")[0].strip()
|
|
|
else:
|
|
|
json_str = response.strip()
|
|
|
|
|
|
# 使用json.loads的strict=False来处理转义字符
|
|
|
# 但这还不够，我们需要更智能的处理
|
|
|
try:
|
|
|
problem_data = json.loads(json_str)
|
|
|
except json.JSONDecodeError as e:
|
|
|
# 如果解析失败，尝试修复常见的LaTeX转义问题
|
|
|
# 方法：先将字符串中的单个反斜杠替换为双反斜杠（但保留已经转义的）
|
|
|
# 这样LaTeX的 \frac 会变成 \\frac，在JSON中是合法的
|
|
|
|
|
|
# 使用正则表达式：找到所有未转义的反斜杠（不是\\的\）
|
|
|
# 并将其替换为\\
|
|
|
fixed_json_str = re.sub(r'(?<!\\)\\(?!["\\/bfnrtu])', r'\\\\', json_str)
|
|
|
|
|
|
try:
|
|
|
problem_data = json.loads(fixed_json_str)
|
|
|
except json.JSONDecodeError:
|
|
|
# 如果还是失败，打印错误信息并抛出
|
|
|
print(f"❌ JSON解析失败:")
|
|
|
print(f"原始响应: {response[:500]}...")
|
|
|
print(f"提取的JSON: {json_str[:500]}...")
|
|
|
raise
|
|
|
|
|
|
# 验证必需字段
|
|
|
if "problem" not in problem_data or "answer" not in problem_data:
|
|
|
raise ValueError("缺少必需字段: problem 或 answer")
|
|
|
|
|
|
# 验证答案范围
|
|
|
answer = int(problem_data.get("answer", 0))
|
|
|
if not (0 <= answer <= 999):
|
|
|
print(f"⚠️ 答案超出范围: {answer}，调整为0-999范围内")
|
|
|
answer = max(0, min(999, answer))
|
|
|
problem_data["answer"] = answer
|
|
|
|
|
|
# 确保有默认值
|
|
|
problem_data.setdefault("solution", "No solution provided")
|
|
|
problem_data.setdefault("topic", "Uncategorized")
|
|
|
|
|
|
return problem_data
|
|
|
|
|
|
def _get_default_problem(self) -> Dict[str, Any]:
|
|
|
"""获取默认题目（生成失败时使用）"""
|
|
|
return {
|
|
|
"problem": "生成失败，请重新生成",
|
|
|
"answer": 0,
|
|
|
"solution": "N/A",
|
|
|
"topic": "未知"
|
|
|
}
|
|
|
|
|
|
def generate_batch(
|
|
|
self,
|
|
|
num_problems: int = 30,
|
|
|
checkpoint_path: str = None
|
|
|
) -> List[Dict[str, Any]]:
|
|
|
"""
|
|
|
批量生成题目
|
|
|
|
|
|
Args:
|
|
|
num_problems: 生成题目数量
|
|
|
checkpoint_path: 检查点文件路径（用于保存进度）
|
|
|
|
|
|
Returns:
|
|
|
题目列表
|
|
|
"""
|
|
|
print(f"\n🎯 开始生成AIME题目")
|
|
|
print(f" 目标数量: {num_problems}")
|
|
|
print(f" 生成模型: {self.llm.model}")
|
|
|
print(f" 延迟设置: {self.delay_seconds}秒/题")
|
|
|
|
|
|
# 尝试从检查点恢复
|
|
|
problems = []
|
|
|
start_index = 0
|
|
|
|
|
|
if checkpoint_path and os.path.exists(checkpoint_path):
|
|
|
print(f"\n📂 发现检查点文件，尝试恢复...")
|
|
|
try:
|
|
|
with open(checkpoint_path, 'r', encoding='utf-8') as f:
|
|
|
problems = json.load(f)
|
|
|
start_index = len(problems)
|
|
|
print(f" ✓ 已恢复 {start_index} 个题目，从第 {start_index + 1} 个继续")
|
|
|
except Exception as e:
|
|
|
print(f" ⚠️ 恢复失败: {e}，从头开始")
|
|
|
problems = []
|
|
|
start_index = 0
|
|
|
|
|
|
# 生成题目（使用tqdm显示进度）
|
|
|
with tqdm(total=num_problems, initial=start_index, desc="生成AIME题目", unit="题") as pbar:
|
|
|
last_call_time = 0 # 上次API调用的时间
|
|
|
|
|
|
for i in range(start_index, num_problems):
|
|
|
# 计算距离上次调用的时间
|
|
|
if last_call_time > 0:
|
|
|
elapsed = time.time() - last_call_time
|
|
|
# 如果距离上次调用不足delay_seconds，则等待
|
|
|
if elapsed < self.delay_seconds:
|
|
|
wait_time = self.delay_seconds - elapsed
|
|
|
tqdm.write(f"⏳ 等待 {wait_time:.1f} 秒以避免速率限制...")
|
|
|
time.sleep(wait_time)
|
|
|
|
|
|
# 记录开始时间
|
|
|
start_time = time.time()
|
|
|
|
|
|
# 生成题目
|
|
|
problem = self.generate_single()
|
|
|
problem["id"] = f"gen_aime_{i + 1}"
|
|
|
problem["generated_at"] = datetime.now().isoformat()
|
|
|
|
|
|
# 记录结束时间
|
|
|
last_call_time = time.time()
|
|
|
generation_time = last_call_time - start_time
|
|
|
|
|
|
problems.append(problem)
|
|
|
|
|
|
# 更新进度条描述
|
|
|
pbar.set_postfix({
|
|
|
"主题": problem.get('topic', 'N/A'),
|
|
|
"答案": problem.get('answer', 'N/A'),
|
|
|
"耗时": f"{generation_time:.1f}s"
|
|
|
})
|
|
|
pbar.update(1)
|
|
|
|
|
|
# 保存检查点
|
|
|
if checkpoint_path:
|
|
|
try:
|
|
|
with open(checkpoint_path, 'w', encoding='utf-8') as f:
|
|
|
json.dump(problems, f, ensure_ascii=False, indent=2)
|
|
|
except Exception as e:
|
|
|
tqdm.write(f"⚠️ 保存检查点失败: {e}")
|
|
|
|
|
|
print(f"\n✅ 生成完成！共 {len(problems)} 个题目")
|
|
|
return problems
|
|
|
|
|
|
def save_problems(
|
|
|
self,
|
|
|
problems: List[Dict[str, Any]],
|
|
|
output_path: str
|
|
|
):
|
|
|
"""保存题目到文件"""
|
|
|
# 确保目录存在
|
|
|
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
|
|
|
|
|
with open(output_path, 'w', encoding='utf-8') as f:
|
|
|
json.dump(problems, f, ensure_ascii=False, indent=2)
|
|
|
|
|
|
print(f"\n💾 题目已保存: {output_path}")
|
|
|
|
|
|
def generate_and_save(
|
|
|
self,
|
|
|
num_problems: int = 30,
|
|
|
output_dir: str = "data_generation/generated_data"
|
|
|
) -> str:
|
|
|
"""生成并保存题目"""
|
|
|
# 创建输出目录
|
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
|
|
# 清理旧的检查点文件
|
|
|
for file in os.listdir(output_dir):
|
|
|
if file.startswith("checkpoint_") and file.endswith(".json"):
|
|
|
old_checkpoint = os.path.join(output_dir, file)
|
|
|
try:
|
|
|
os.remove(old_checkpoint)
|
|
|
print(f"🗑️ 已删除旧检查点文件: {file}")
|
|
|
except Exception as e:
|
|
|
print(f"⚠️ 删除旧检查点失败: {e}")
|
|
|
|
|
|
# 设置检查点路径
|
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
|
checkpoint_path = os.path.join(output_dir, f"checkpoint_{timestamp}.json")
|
|
|
|
|
|
# 生成题目（带检查点）
|
|
|
problems = self.generate_batch(num_problems, checkpoint_path=checkpoint_path)
|
|
|
|
|
|
# 保存题目
|
|
|
output_path = os.path.join(output_dir, f"aime_generated_{timestamp}.json")
|
|
|
self.save_problems(problems, output_path)
|
|
|
|
|
|
# 生成统计报告
|
|
|
self._generate_statistics_report(problems, output_dir, timestamp)
|
|
|
|
|
|
# 删除检查点文件
|
|
|
if os.path.exists(checkpoint_path):
|
|
|
try:
|
|
|
os.remove(checkpoint_path)
|
|
|
print(f"\n🗑️ 已删除检查点文件")
|
|
|
except Exception as e:
|
|
|
print(f"\n⚠️ 删除检查点文件失败: {e}")
|
|
|
|
|
|
return output_path
|
|
|
|
|
|
def _generate_statistics_report(
|
|
|
self,
|
|
|
problems: List[Dict[str, Any]],
|
|
|
output_dir: str,
|
|
|
timestamp: str
|
|
|
):
|
|
|
"""生成统计报告"""
|
|
|
# 统计主题分布
|
|
|
topics = {}
|
|
|
answers = []
|
|
|
|
|
|
for problem in problems:
|
|
|
topic = problem.get("topic", "未知")
|
|
|
topics[topic] = topics.get(topic, 0) + 1
|
|
|
|
|
|
if "answer" in problem:
|
|
|
answers.append(problem["answer"])
|
|
|
|
|
|
# 生成报告
|
|
|
report = f"""# AIME题目生成统计报告
|
|
|
|
|
|
## 基本信息
|
|
|
|
|
|
- **生成时间**: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
|
|
|
- **题目数量**: {len(problems)}
|
|
|
|
|
|
## 主题分布
|
|
|
|
|
|
| 主题 | 数量 | 占比 |
|
|
|
|------|------|------|
|
|
|
"""
|
|
|
|
|
|
for topic, count in sorted(topics.items(), key=lambda x: x[1], reverse=True):
|
|
|
percentage = count / len(problems) * 100
|
|
|
report += f"| {topic} | {count} | {percentage:.1f}% |\n"
|
|
|
|
|
|
if answers:
|
|
|
report += f"""
|
|
|
## 答案分析
|
|
|
|
|
|
- **平均答案**: {sum(answers) / len(answers):.2f}
|
|
|
- **最小答案**: {min(answers)}
|
|
|
- **最大答案**: {max(answers)}
|
|
|
- **答案范围**: {min(answers)}-{max(answers)}
|
|
|
"""
|
|
|
|
|
|
report += f"""
|
|
|
## 题目列表
|
|
|
|
|
|
| ID | 主题 | 答案 |
|
|
|
|-----|------|------|
|
|
|
"""
|
|
|
|
|
|
for problem in problems[:10]: # 只显示前10个
|
|
|
report += f"| {problem.get('id', 'N/A')} | {problem.get('topic', 'N/A')} | {problem.get('answer', 'N/A')} |\n"
|
|
|
|
|
|
if len(problems) > 10:
|
|
|
report += f"\n*（仅显示前10个题目，完整列表请查看JSON文件）*\n"
|
|
|
|
|
|
report += f"""
|
|
|
---
|
|
|
|
|
|
*报告生成时间: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}*
|
|
|
"""
|
|
|
|
|
|
# 保存报告
|
|
|
report_path = os.path.join(output_dir, f"generation_report_{timestamp}.md")
|
|
|
with open(report_path, 'w', encoding='utf-8') as f:
|
|
|
f.write(report)
|
|
|
|
|
|
print(f"📊 统计报告已保存: {report_path}")
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
# 创建生成器
|
|
|
generator = AIMEGenerator()
|
|
|
|
|
|
# 生成30个题目
|
|
|
output_path = generator.generate_and_save(num_problems=30)
|
|
|
|
|
|
print(f"\n✅ 完成！生成的题目保存在: {output_path}")
|
|
|
|