diff --git a/scripts/eval_toolcall.py b/scripts/eval_toolcall.py index da078b4..d159eff 100644 --- a/scripts/eval_toolcall.py +++ b/scripts/eval_toolcall.py @@ -12,7 +12,7 @@ from datetime import datetime from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer from openai import OpenAI from model.model_minimind import MiniMindConfig, MiniMindForCausalLM -from trainer.trainer_utils import setup_seed, get_model_params +from trainer.trainer_utils import setup_seed, get_model_params, safe_math_eval warnings.filterwarnings('ignore') TOOLS = [ @@ -27,7 +27,7 @@ TOOLS = [ ] MOCK_RESULTS = { - "calculate_math": lambda args: {"result": str(eval(str(args.get("expression", "0")).replace("^", "**").replace("×", "*").replace("÷", "/").replace("−", "-").replace("²", "**2").replace("³", "**3").replace("（", "(").replace("）", ")")))}, + "calculate_math": lambda args: {"result": str(safe_math_eval(args.get("expression", "0")))}, "get_current_time": lambda args: {"datetime": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "timezone": args.get("timezone", "Asia/Shanghai")}, "random_number": lambda args: {"result": random.randint(int(args.get("min", 0)), int(args.get("max", 100)))}, "text_length": lambda args: {"characters": len(args.get("text", "")), "words": len(args.get("text", "").split())}, diff --git a/scripts/web_demo.py b/scripts/web_demo.py index 7debdcc..d07d96d 100644 --- a/scripts/web_demo.py +++ b/scripts/web_demo.py @@ -2,6 +2,7 @@ import random import re import json import os +import sys from threading import Thread import torch @@ -9,6 +10,9 @@ import numpy as np import streamlit as st from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) +from trainer.trainer_utils import safe_math_eval + st.set_page_config(page_title="MiniMind", initial_sidebar_state="collapsed") st.markdown(""" @@ -125,7 +129,7 @@ def execute_tool(tool_name, args): import datetime try: if tool_name == 'calculate_math': - return {"result": eval(args.get('expression', '0'))} + return {"result": safe_math_eval(args.get('expression', '0'))} elif tool_name == 'get_current_time': tz = args.get('timezone', 'Asia/Shanghai') return {"result": datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')} diff --git a/trainer/train_agent.py b/trainer/train_agent.py index 6b4360f..50520aa 100644 --- a/trainer/train_agent.py +++ b/trainer/train_agent.py @@ -10,7 +10,6 @@ import gc import json import math import random -import signal import argparse import warnings import torch @@ -24,7 +23,7 @@ from torch.optim.lr_scheduler import CosineAnnealingLR from transformers import AutoTokenizer from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from dataset.lm_dataset import AgentRLDataset -from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel +from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel, safe_math_eval from trainer.rollout_engine import create_rollout_engine, compute_per_token_logps warnings.filterwarnings('ignore') @@ -55,7 +54,7 @@ UNIT_DATA = {"km_miles": 0.621371, "miles_km": 1.60934, "kg_pounds": 2.20462, "p # ======== 模拟执行 ======== MOCK_RESULTS = { - "calculate_math": lambda args: {"result": str(eval(str(args.get("expression", "0")).replace("^", "**").replace("×", "*").replace("÷", "/").replace("−", "-").replace("（", "(").replace("）", ")"), {"__builtins__": {}, "math": math}))}, + "calculate_math": lambda args: {"result": str(safe_math_eval(args.get("expression", "0")))}, "unit_converter": lambda args: {"result": round(float(args.get("value", 0)) * UNIT_DATA.get(f"{args.get('from_unit', '').lower()}_{args.get('to_unit', '').lower()}", 1), 4)}, "get_current_weather": lambda args: (lambda w: {"city": args.get("location"), "temperature": w[0], "humidity": "65%", "condition": w[1]})(WEATHER_DATA.get(args.get("location"), ("22°C", "晴"))), "get_current_time": lambda args: {"datetime": TIME_DATA.get(args.get("timezone", "Asia/Shanghai"), "2025-03-07 14:30:00"), "timezone": args.get("timezone", "Asia/Shanghai")}, @@ -85,14 +84,9 @@ def execute_tool(name, args): fn = MOCK_RESULTS.get(name) if not fn: return None try: - signal.signal(signal.SIGALRM, lambda *_: (_ for _ in ()).throw(TimeoutError())) - signal.alarm(1) return fn(args) - except: + except Exception: return None - finally: - try: signal.alarm(0) - except: pass # ======== 多轮 Rollout ======== def rollout_single(rollout_engine, tokenizer, messages, tools, max_turns=3, max_new_tokens=256, thinking_ratio=0.5, device="cuda"): diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py index b6211d5..5ebf727 100644 --- a/trainer/trainer_utils.py +++ b/trainer/trainer_utils.py @@ -5,6 +5,8 @@ import os import sys __package__ = "trainer" sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) +import ast +import operator import random import math import numpy as np @@ -174,4 +176,53 @@ class LMForRewardModel: {"role": "assistant", "content": response} ] score = self.model.get_score(self.tokenizer, eval_messages) - return max(min(score, 3.0), -3.0) \ No newline at end of file + return max(min(score, 3.0), -3.0) + + +# ===== 数学表达式安全求值（替代 eval，避免执行模型生成的任意代码） ===== +_MATH_FUNCS = {name: getattr(math, name) for name in ( + 'sqrt', 'pow', 'exp', 'log', 'log2', 'log10', 'sin', 'cos', 'tan', 'asin', 'acos', 'atan', 'atan2', + 'sinh', 'cosh', 'tanh', 'floor', 'ceil', 'trunc', 'fabs', 'fmod', 'hypot', 'gcd', 'degrees', 'radians' +)} +_MATH_CONSTS = {'pi': math.pi, 'e': math.e, 'tau': math.tau} +_BIN_OPS = {ast.Add: operator.add, ast.Sub: operator.sub, ast.Mult: operator.mul, ast.Div: operator.truediv, + ast.FloorDiv: operator.floordiv, ast.Mod: operator.mod, ast.Pow: operator.pow} +_UNARY_OPS = {ast.UAdd: operator.pos, ast.USub: operator.neg} +_MAX_EXPR_LEN, _MAX_EXP, _MAX_DIGITS = 512, 1e4, 100 # 表达式长度 / 指数 / 结果位数上限 + + +def _check_pow(base, exp): # 防止 9**9**9 这类表达式把进程算死 + if abs(exp) > _MAX_EXP or (abs(base) > 1 and abs(exp) * math.log10(abs(base)) > _MAX_DIGITS): + raise ValueError('幂运算结果过大') + + +def _resolve_name(node): # 同时兼容 sqrt(4) 与 math.sqrt(4) 两种写法 + if isinstance(node, ast.Name): return node.id + if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name) and node.value.id == 'math': return node.attr + return None + + +def _eval_ast(node): + if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)) and not isinstance(node.value, bool): + return node.value + if isinstance(node, (ast.Name, ast.Attribute)) and _resolve_name(node) in _MATH_CONSTS: + return _MATH_CONSTS[_resolve_name(node)] + if isinstance(node, ast.UnaryOp) and type(node.op) in _UNARY_OPS: + return _UNARY_OPS[type(node.op)](_eval_ast(node.operand)) + if isinstance(node, ast.BinOp) and type(node.op) in _BIN_OPS: + left, right = _eval_ast(node.left), _eval_ast(node.right) + if isinstance(node.op, ast.Pow): _check_pow(left, right) + return _BIN_OPS[type(node.op)](left, right) + if isinstance(node, ast.Call) and _resolve_name(node.func) in _MATH_FUNCS and not node.keywords: + name, args = _resolve_name(node.func), [_eval_ast(arg) for arg in node.args] + if name == 'pow' and len(args) == 2: _check_pow(*args) + return _MATH_FUNCS[name](*args) + raise ValueError(f'不支持的表达式语法: {type(node).__name__}') + + +def safe_math_eval(expression): + """对模型生成的数学表达式求值：只放行算术运算与 math 白名单函数，不执行任意代码。""" + expr = str(expression).replace('^', '**').replace('×', '*').replace('÷', '/').replace('−', '-') + expr = expr.replace('²', '**2').replace('³', '**3').replace('（', '(').replace('）', ')').strip() + if not expr or len(expr) > _MAX_EXPR_LEN: raise ValueError('表达式为空或过长') + return _eval_ast(ast.parse(expr, mode='eval').body)