agentic
baihaochen
parent
5012b232ff
commit
3233d9872c
|
|
@ -0,0 +1,21 @@
|
|||
from .data_analysis_env import (
|
||||
AGENTIC_SYSTEM_PROMPT,
|
||||
AgenticToolEnv,
|
||||
average_agentic_metrics,
|
||||
extract_final_answer,
|
||||
format_agentic_user_prompt,
|
||||
get_agentic_tools,
|
||||
parse_tool_calls,
|
||||
score_agentic_trajectory,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AGENTIC_SYSTEM_PROMPT",
|
||||
"AgenticToolEnv",
|
||||
"average_agentic_metrics",
|
||||
"extract_final_answer",
|
||||
"format_agentic_user_prompt",
|
||||
"get_agentic_tools",
|
||||
"parse_tool_calls",
|
||||
"score_agentic_trajectory",
|
||||
]
|
||||
|
|
@ -0,0 +1,566 @@
|
|||
import ast
|
||||
import contextlib
|
||||
import io
|
||||
import json
|
||||
import math
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import re
|
||||
import statistics
|
||||
from collections import Counter
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
|
||||
AGENTIC_SYSTEM_PROMPT = (
|
||||
"你是一个面向运营数据分析的 Agent。你需要根据用户问题选择合适工具，"
|
||||
"优先检索指标口径，再用 Python 对给定 CSV 数据进行计算；最终答案必须简洁，"
|
||||
"包含关键结论和必要数值。工具调用必须使用 <tool_call> 包裹的 JSON。"
|
||||
)
|
||||
|
||||
|
||||
AGENTIC_TOOL_SCHEMAS = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "retriever",
|
||||
"description": "检索指标口径、表字段说明或业务规则，返回最相关的知识片段。",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {"type": "string", "description": "检索查询"},
|
||||
"top_k": {"type": "integer", "description": "返回片段数量", "default": 3},
|
||||
},
|
||||
"required": ["query"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "python_executor",
|
||||
"description": "在受限 Python 环境中使用 pandas/numpy 读取给定 CSV 并完成数据分析。",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"code": {
|
||||
"type": "string",
|
||||
"description": "Python 代码。把最终结果赋值给 result 或 answer。",
|
||||
}
|
||||
},
|
||||
"required": ["code"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "calculator",
|
||||
"description": "执行确定性数学表达式计算，适合简单四则运算、比例和幂运算。",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"expression": {"type": "string", "description": "数学表达式，如 1280 * 0.15 * 3.2"}
|
||||
},
|
||||
"required": ["expression"],
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def get_agentic_tools(names=None):
|
||||
if names is None:
|
||||
names = ["retriever", "python_executor", "calculator"]
|
||||
name_set = set(names)
|
||||
return [tool for tool in AGENTIC_TOOL_SCHEMAS if tool["function"]["name"] in name_set]
|
||||
|
||||
|
||||
def json_dumps(data):
|
||||
return json.dumps(data, ensure_ascii=False, default=str)
|
||||
|
||||
|
||||
def format_agentic_user_prompt(sample):
|
||||
table_lines = []
|
||||
for table in sample.get("tables", []):
|
||||
table_lines.append(
|
||||
f"- {table.get('name', '')}: path={table.get('path', '')}; {table.get('description', '')}"
|
||||
)
|
||||
tool_names = ", ".join(sample.get("tools", ["retriever", "python_executor", "calculator"]))
|
||||
return (
|
||||
f"用户问题：{sample.get('question', '')}\n\n"
|
||||
f"可用数据表：\n" + "\n".join(table_lines) + "\n\n"
|
||||
f"可用工具：{tool_names}\n"
|
||||
"要求：先确认指标口径，再调用工具计算；最终答案只输出结论，不要编造数据。"
|
||||
)
|
||||
|
||||
|
||||
def parse_tool_calls(text):
|
||||
calls = []
|
||||
for match in re.findall(r"<tool_call>(.*?)</tool_call>", text or "", re.DOTALL):
|
||||
raw = match.strip()
|
||||
try:
|
||||
data = json.loads(raw)
|
||||
args = data.get("arguments", {})
|
||||
if isinstance(args, str):
|
||||
try:
|
||||
args = json.loads(args)
|
||||
except Exception:
|
||||
args = {}
|
||||
calls.append(
|
||||
{
|
||||
"name": data.get("name", ""),
|
||||
"arguments": args if isinstance(args, dict) else {},
|
||||
"raw": raw,
|
||||
"valid_json": True,
|
||||
}
|
||||
)
|
||||
except Exception:
|
||||
calls.append({"name": "", "arguments": {}, "raw": raw, "valid_json": False})
|
||||
return calls
|
||||
|
||||
|
||||
def strip_thinking(text):
|
||||
text = text or ""
|
||||
return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()
|
||||
|
||||
|
||||
def extract_final_answer(turn_outputs):
|
||||
if isinstance(turn_outputs, str):
|
||||
turn_outputs = [turn_outputs]
|
||||
if not turn_outputs:
|
||||
return ""
|
||||
text = strip_thinking(turn_outputs[-1])
|
||||
if "</tool_call>" in text:
|
||||
text = text.split("</tool_call>")[-1]
|
||||
text = re.sub(r"<tool_call>.*?</tool_call>", "", text, flags=re.DOTALL)
|
||||
return text.strip()
|
||||
|
||||
|
||||
def _tokenize(text):
|
||||
return re.findall(r"[\w\u4e00-\u9fff]+", str(text).lower())
|
||||
|
||||
|
||||
def _overlap_score(query, text):
|
||||
q = Counter(_tokenize(query))
|
||||
d = Counter(_tokenize(text))
|
||||
return sum(min(q[token], d[token]) for token in q)
|
||||
|
||||
|
||||
def _safe_number(value, default=0.0):
|
||||
try:
|
||||
return float(value)
|
||||
except Exception:
|
||||
return default
|
||||
|
||||
|
||||
class _SafeArithmetic(ast.NodeVisitor):
|
||||
allowed_binops = {
|
||||
ast.Add: lambda a, b: a + b,
|
||||
ast.Sub: lambda a, b: a - b,
|
||||
ast.Mult: lambda a, b: a * b,
|
||||
ast.Div: lambda a, b: a / b,
|
||||
ast.FloorDiv: lambda a, b: a // b,
|
||||
ast.Mod: lambda a, b: a % b,
|
||||
ast.Pow: lambda a, b: a ** b,
|
||||
}
|
||||
allowed_unary = {
|
||||
ast.UAdd: lambda a: +a,
|
||||
ast.USub: lambda a: -a,
|
||||
}
|
||||
allowed_funcs = {
|
||||
"abs": abs,
|
||||
"round": round,
|
||||
"sqrt": math.sqrt,
|
||||
"log": math.log,
|
||||
"log10": math.log10,
|
||||
"ceil": math.ceil,
|
||||
"floor": math.floor,
|
||||
"pow": pow,
|
||||
}
|
||||
|
||||
def visit_Expression(self, node):
|
||||
return self.visit(node.body)
|
||||
|
||||
def visit_Constant(self, node):
|
||||
if isinstance(node.value, (int, float)):
|
||||
return float(node.value)
|
||||
raise ValueError("calculator only supports numeric constants")
|
||||
|
||||
def visit_BinOp(self, node):
|
||||
op = type(node.op)
|
||||
if op not in self.allowed_binops:
|
||||
raise ValueError(f"operator not allowed: {op.__name__}")
|
||||
return self.allowed_binops[op](self.visit(node.left), self.visit(node.right))
|
||||
|
||||
def visit_UnaryOp(self, node):
|
||||
op = type(node.op)
|
||||
if op not in self.allowed_unary:
|
||||
raise ValueError(f"operator not allowed: {op.__name__}")
|
||||
return self.allowed_unary[op](self.visit(node.operand))
|
||||
|
||||
def visit_Call(self, node):
|
||||
if not isinstance(node.func, ast.Name) or node.func.id not in self.allowed_funcs:
|
||||
raise ValueError("function not allowed")
|
||||
args = [self.visit(arg) for arg in node.args]
|
||||
return self.allowed_funcs[node.func.id](*args)
|
||||
|
||||
def generic_visit(self, node):
|
||||
raise ValueError(f"node not allowed: {type(node).__name__}")
|
||||
|
||||
|
||||
def safe_calculate(expression):
|
||||
normalized = (
|
||||
str(expression)
|
||||
.replace("^", "**")
|
||||
.replace("×", "*")
|
||||
.replace("÷", "/")
|
||||
.replace("−", "-")
|
||||
.replace("（", "(")
|
||||
.replace("）", ")")
|
||||
)
|
||||
tree = ast.parse(normalized, mode="eval")
|
||||
return _SafeArithmetic().visit(tree)
|
||||
|
||||
|
||||
def _jsonable(value):
|
||||
if isinstance(value, pd.DataFrame):
|
||||
return value.to_dict("records")
|
||||
if isinstance(value, pd.Series):
|
||||
return value.to_dict()
|
||||
if isinstance(value, np.generic):
|
||||
return value.item()
|
||||
if isinstance(value, dict):
|
||||
return {str(k): _jsonable(v) for k, v in value.items()}
|
||||
if isinstance(value, (list, tuple, set)):
|
||||
return [_jsonable(v) for v in value]
|
||||
return value
|
||||
|
||||
|
||||
def _validate_python_ast(code):
|
||||
banned_nodes = (ast.Import, ast.ImportFrom, ast.Global, ast.Nonlocal, ast.Lambda, ast.ClassDef)
|
||||
banned_names = {
|
||||
"__import__",
|
||||
"eval",
|
||||
"exec",
|
||||
"compile",
|
||||
"open",
|
||||
"input",
|
||||
"globals",
|
||||
"locals",
|
||||
"vars",
|
||||
"dir",
|
||||
"getattr",
|
||||
"setattr",
|
||||
"delattr",
|
||||
"os",
|
||||
"sys",
|
||||
"subprocess",
|
||||
"socket",
|
||||
"requests",
|
||||
"urllib",
|
||||
"pathlib",
|
||||
"shutil",
|
||||
}
|
||||
banned_attrs = {
|
||||
"to_csv",
|
||||
"to_excel",
|
||||
"to_pickle",
|
||||
"to_parquet",
|
||||
"write",
|
||||
"remove",
|
||||
"unlink",
|
||||
"rmdir",
|
||||
"mkdir",
|
||||
"rename",
|
||||
"replace",
|
||||
"system",
|
||||
"popen",
|
||||
}
|
||||
tree = ast.parse(code, mode="exec")
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, banned_nodes):
|
||||
raise ValueError(f"python_executor does not allow {type(node).__name__}")
|
||||
if isinstance(node, ast.Name) and node.id in banned_names:
|
||||
raise ValueError(f"name not allowed: {node.id}")
|
||||
if isinstance(node, ast.Attribute) and node.attr in banned_attrs:
|
||||
raise ValueError(f"attribute not allowed: {node.attr}")
|
||||
return tree
|
||||
|
||||
|
||||
class _PandasProxy:
|
||||
def __init__(self, repo_root, allowed_paths):
|
||||
self.repo_root = repo_root
|
||||
self.allowed_paths = {os.path.abspath(path) for path in allowed_paths}
|
||||
|
||||
def read_csv(self, path, *args, **kwargs):
|
||||
candidates = []
|
||||
raw = str(path)
|
||||
if os.path.isabs(raw):
|
||||
candidates.append(os.path.abspath(raw))
|
||||
candidates.append(os.path.abspath(raw))
|
||||
candidates.append(os.path.abspath(os.path.join(self.repo_root, raw)))
|
||||
allowed_by_basename = {os.path.basename(path): path for path in self.allowed_paths}
|
||||
if raw in allowed_by_basename:
|
||||
candidates.append(allowed_by_basename[raw])
|
||||
for candidate in candidates:
|
||||
if candidate in self.allowed_paths:
|
||||
return pd.read_csv(candidate, *args, **kwargs)
|
||||
raise PermissionError(f"read_csv path is not in sample tables: {raw}")
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(pd, name)
|
||||
|
||||
|
||||
def _python_worker(code, repo_root, allowed_paths, queue):
|
||||
try:
|
||||
tree = _validate_python_ast(code)
|
||||
output = io.StringIO()
|
||||
safe_builtins = {
|
||||
"abs": abs,
|
||||
"all": all,
|
||||
"any": any,
|
||||
"bool": bool,
|
||||
"dict": dict,
|
||||
"enumerate": enumerate,
|
||||
"float": float,
|
||||
"int": int,
|
||||
"len": len,
|
||||
"list": list,
|
||||
"max": max,
|
||||
"min": min,
|
||||
"pow": pow,
|
||||
"print": print,
|
||||
"range": range,
|
||||
"round": round,
|
||||
"set": set,
|
||||
"sorted": sorted,
|
||||
"str": str,
|
||||
"sum": sum,
|
||||
"tuple": tuple,
|
||||
"zip": zip,
|
||||
}
|
||||
local_vars = {}
|
||||
global_vars = {
|
||||
"__builtins__": safe_builtins,
|
||||
"pd": _PandasProxy(repo_root, allowed_paths),
|
||||
"np": np,
|
||||
"math": math,
|
||||
"json": json,
|
||||
"statistics": statistics,
|
||||
}
|
||||
with contextlib.redirect_stdout(output):
|
||||
exec(compile(tree, "<agentic_python_executor>", "exec"), global_vars, local_vars)
|
||||
value = local_vars.get("result", local_vars.get("answer", output.getvalue().strip()))
|
||||
queue.put({"ok": True, "result": _jsonable(value), "error": "", "artifacts": []})
|
||||
except Exception as exc:
|
||||
queue.put({"ok": False, "result": "", "error": str(exc)[:500], "artifacts": []})
|
||||
|
||||
|
||||
class AgenticToolEnv:
|
||||
def __init__(self, sample, repo_root=None, timeout=3):
|
||||
self.sample = sample
|
||||
self.repo_root = os.path.abspath(repo_root or os.getcwd())
|
||||
self.timeout = timeout
|
||||
self.allowed_paths = [self._resolve_table_path(table.get("path", "")) for table in sample.get("tables", [])]
|
||||
|
||||
def _resolve_table_path(self, path):
|
||||
if os.path.isabs(path):
|
||||
return os.path.abspath(path)
|
||||
return os.path.abspath(os.path.join(self.repo_root, path))
|
||||
|
||||
def execute(self, name, arguments=None):
|
||||
arguments = arguments or {}
|
||||
if name == "retriever":
|
||||
return self._retriever(arguments)
|
||||
if name == "calculator":
|
||||
return self._calculator(arguments)
|
||||
if name == "python_executor":
|
||||
return self._python_executor(arguments)
|
||||
return {"ok": False, "result": "", "error": f"unknown tool: {name}", "artifacts": []}
|
||||
|
||||
def _retriever(self, arguments):
|
||||
query = str(arguments.get("query", self.sample.get("question", "")))
|
||||
top_k = max(1, min(int(arguments.get("top_k", 3) or 3), 8))
|
||||
docs = self.sample.get("documents", [])
|
||||
scored = []
|
||||
for doc in docs:
|
||||
text = f"{doc.get('title', '')} {doc.get('text', '')}"
|
||||
scored.append((_overlap_score(query, text), doc))
|
||||
scored.sort(key=lambda item: item[0], reverse=True)
|
||||
result = []
|
||||
for score, doc in scored[:top_k]:
|
||||
result.append(
|
||||
{
|
||||
"id": doc.get("id", ""),
|
||||
"title": doc.get("title", ""),
|
||||
"text": doc.get("text", ""),
|
||||
"score": score,
|
||||
}
|
||||
)
|
||||
return {"ok": True, "result": result, "error": "", "artifacts": []}
|
||||
|
||||
def _calculator(self, arguments):
|
||||
expression = arguments.get("expression", "")
|
||||
try:
|
||||
value = safe_calculate(expression)
|
||||
except Exception as exc:
|
||||
return {"ok": False, "result": "", "error": str(exc)[:500], "artifacts": []}
|
||||
return {"ok": True, "result": value, "error": "", "artifacts": []}
|
||||
|
||||
def _python_executor(self, arguments):
|
||||
code = str(arguments.get("code", ""))
|
||||
if not code.strip():
|
||||
return {"ok": False, "result": "", "error": "empty code", "artifacts": []}
|
||||
queue = mp.Queue(maxsize=1)
|
||||
proc = mp.Process(target=_python_worker, args=(code, self.repo_root, self.allowed_paths, queue))
|
||||
proc.start()
|
||||
proc.join(self.timeout)
|
||||
if proc.is_alive():
|
||||
proc.terminate()
|
||||
proc.join()
|
||||
return {"ok": False, "result": "", "error": f"python_executor timeout>{self.timeout}s", "artifacts": []}
|
||||
if queue.empty():
|
||||
return {"ok": False, "result": "", "error": "python_executor returned no result", "artifacts": []}
|
||||
return queue.get()
|
||||
|
||||
|
||||
def validate_tool_call(call, available_names):
|
||||
if not call.get("valid_json", False):
|
||||
return False
|
||||
name = call.get("name", "")
|
||||
args = call.get("arguments", {})
|
||||
if name not in available_names:
|
||||
return False
|
||||
if name == "retriever":
|
||||
return bool(args.get("query"))
|
||||
if name == "calculator":
|
||||
return bool(args.get("expression"))
|
||||
if name == "python_executor":
|
||||
return bool(args.get("code"))
|
||||
return False
|
||||
|
||||
|
||||
def _f1(pred_items, gold_items):
|
||||
pred = set(pred_items)
|
||||
gold = set(gold_items)
|
||||
if not pred and not gold:
|
||||
return 1.0
|
||||
if not pred or not gold:
|
||||
return 0.0
|
||||
inter = len(pred & gold)
|
||||
if inter == 0:
|
||||
return 0.0
|
||||
precision = inter / len(pred)
|
||||
recall = inter / len(gold)
|
||||
return 2 * precision * recall / (precision + recall)
|
||||
|
||||
|
||||
def _numbers(text):
|
||||
return [float(x.replace(",", "")) for x in re.findall(r"[-+]?\d+(?:,\d{3})*(?:\.\d+)?", str(text))]
|
||||
|
||||
|
||||
def score_answer(final_answer, sample):
|
||||
checks = sample.get("checks", {}) or {}
|
||||
score_parts = []
|
||||
contains = checks.get("contains") or []
|
||||
if contains:
|
||||
score_parts.append(float(all(str(item).lower() in final_answer.lower() for item in contains)))
|
||||
if "number" in checks:
|
||||
target = _safe_number(checks.get("number"))
|
||||
tolerance = _safe_number(checks.get("tolerance"), 1e-3)
|
||||
nums = _numbers(final_answer)
|
||||
score_parts.append(float(any(abs(num - target) <= tolerance for num in nums)))
|
||||
if not score_parts and sample.get("answer"):
|
||||
score_parts.append(float(str(sample["answer"]).lower() in final_answer.lower()))
|
||||
return sum(score_parts) / len(score_parts) if score_parts else 0.0
|
||||
|
||||
|
||||
def repetition_penalty(text, n=3, cap=0.5):
|
||||
toks = _tokenize(text)
|
||||
grams = [tuple(toks[i:i + n]) for i in range(len(toks) - n + 1)]
|
||||
return min(cap, (len(grams) - len(set(grams))) * cap * 2 / len(grams)) if grams else 0.0
|
||||
|
||||
|
||||
def score_agentic_trajectory(turn_outputs, sample, repo_root=None, unfinished=False, execute_tools=True):
|
||||
if isinstance(turn_outputs, str):
|
||||
turn_outputs = [turn_outputs]
|
||||
turn_outputs = turn_outputs or []
|
||||
available_names = set(sample.get("tools", ["retriever", "python_executor", "calculator"]))
|
||||
expected_tools = sample.get("expected_tools", [])
|
||||
env = AgenticToolEnv(sample, repo_root=repo_root)
|
||||
|
||||
calls = []
|
||||
tag_errors = 0
|
||||
for turn in turn_outputs:
|
||||
tag_errors += abs(str(turn).count("<tool_call>") - str(turn).count("</tool_call>"))
|
||||
calls.extend(parse_tool_calls(turn))
|
||||
|
||||
valid_flags = [validate_tool_call(call, available_names) for call in calls]
|
||||
valid_count = sum(valid_flags)
|
||||
invalid_count = len(calls) - valid_count
|
||||
selected_tools = [call.get("name", "") for call, ok in zip(calls, valid_flags) if ok]
|
||||
tool_f1 = _f1(selected_tools, expected_tools)
|
||||
|
||||
exec_success = 1.0
|
||||
if execute_tools and calls:
|
||||
successes = []
|
||||
for call, ok in zip(calls, valid_flags):
|
||||
if not ok:
|
||||
successes.append(False)
|
||||
continue
|
||||
result = env.execute(call["name"], call["arguments"])
|
||||
successes.append(bool(result.get("ok")))
|
||||
exec_success = sum(successes) / len(successes) if successes else 0.0
|
||||
|
||||
final_answer = extract_final_answer(turn_outputs)
|
||||
answer_score = 0.0 if unfinished else score_answer(final_answer, sample)
|
||||
task_success = float(answer_score >= 0.999)
|
||||
schema_valid_rate = valid_count / len(calls) if calls else 0.0
|
||||
invalid_action_rate = invalid_count / len(calls) if calls else 0.0
|
||||
format_score = 1.0 if tag_errors == 0 and all(call.get("valid_json") for call in calls) else 0.0
|
||||
turn_count = len(turn_outputs)
|
||||
expected_turns = max(1, int(sample.get("max_turns", 4)))
|
||||
turn_efficiency = max(0.0, 1.0 - max(0, turn_count - expected_turns) / expected_turns)
|
||||
repeat_penalty = repetition_penalty(final_answer)
|
||||
|
||||
reward = (
|
||||
2.5 * task_success
|
||||
+ 1.2 * answer_score
|
||||
+ 0.8 * tool_f1
|
||||
+ 0.6 * schema_valid_rate
|
||||
+ 0.6 * exec_success
|
||||
+ 0.4 * format_score
|
||||
+ 0.2 * turn_efficiency
|
||||
- 0.8 * invalid_action_rate
|
||||
- 0.3 * tag_errors
|
||||
- repeat_penalty
|
||||
- (0.5 if unfinished else 0.0)
|
||||
)
|
||||
reward = max(min(reward, 5.0), -3.0)
|
||||
return reward, {
|
||||
"reward_total": reward,
|
||||
"task_success": task_success,
|
||||
"answer_score": answer_score,
|
||||
"tool_selection_f1": tool_f1,
|
||||
"schema_valid_rate": schema_valid_rate,
|
||||
"exec_success_rate": exec_success,
|
||||
"invalid_action_rate": invalid_action_rate,
|
||||
"format_score": format_score,
|
||||
"turn_efficiency": turn_efficiency,
|
||||
"repetition_penalty": repeat_penalty,
|
||||
"num_tool_calls": len(calls),
|
||||
"avg_turns": turn_count,
|
||||
}
|
||||
|
||||
|
||||
def average_agentic_metrics(metrics):
|
||||
if not metrics:
|
||||
return {}
|
||||
keys = sorted({key for row in metrics for key in row})
|
||||
return {
|
||||
key: sum(float(row.get(key, 0.0)) for row in metrics) / len(metrics)
|
||||
for key in keys
|
||||
}
|
||||
|
|
@ -0,0 +1,40 @@
|
|||
import json
|
||||
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
from agentic.data_analysis_env import AGENTIC_SYSTEM_PROMPT, format_agentic_user_prompt, get_agentic_tools
|
||||
|
||||
|
||||
def load_agentic_jsonl(path):
|
||||
rows = []
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if line:
|
||||
rows.append(json.loads(line))
|
||||
return rows
|
||||
|
||||
|
||||
class AgenticRLDataset(Dataset):
|
||||
def __init__(self, jsonl_path, tokenizer=None, max_length=2048):
|
||||
super().__init__()
|
||||
self.samples = load_agentic_jsonl(jsonl_path)
|
||||
self.tokenizer = tokenizer
|
||||
self.max_length = max_length
|
||||
|
||||
def __len__(self):
|
||||
return len(self.samples)
|
||||
|
||||
def __getitem__(self, index):
|
||||
sample = self.samples[index]
|
||||
tools = get_agentic_tools(sample.get("tools"))
|
||||
messages = [
|
||||
{"role": "system", "content": AGENTIC_SYSTEM_PROMPT},
|
||||
{"role": "user", "content": format_agentic_user_prompt(sample)},
|
||||
]
|
||||
return {
|
||||
"id": sample.get("id", str(index)),
|
||||
"messages": messages,
|
||||
"tools": tools,
|
||||
"sample": sample,
|
||||
}
|
||||
|
|
@ -0,0 +1,183 @@
|
|||
# Agentic DataAnalysis：面向运营数据分析的 Agentic RL 系统
|
||||
|
||||
## 项目定位
|
||||
|
||||
面向运营数据分析场景，构建一个 300M 级 Agentic LLM：模型需要先理解指标口径，再调用 `retriever`、`python_executor`、`calculator` 等工具完成多步分析，最终输出可验证的业务结论。
|
||||
|
||||
这条线不重新预训练，直接使用已有 `SearchLM-300M` 预训练权重：
|
||||
|
||||
```text
|
||||
out/pretrain_searchlm_1024_mla.pth
|
||||
```
|
||||
|
||||
后续主流程：
|
||||
|
||||
```text
|
||||
Agent 轨迹 SFT -> Agentic GRPO/CISPO -> 离线评测 -> Demo 展示
|
||||
```
|
||||
|
||||
## 数据构造
|
||||
|
||||
生成运营数据表、指标文档、任务、标准答案和专家工具轨迹：
|
||||
|
||||
```bash
|
||||
python scripts/build_agentic_data_tasks.py \
|
||||
--output_dir dataset/agentic_data \
|
||||
--train_size 30000 \
|
||||
--dev_size 1000 \
|
||||
--test_size 1000
|
||||
```
|
||||
|
||||
输出：
|
||||
|
||||
- `dataset/agentic_data/tables/orders.csv`
|
||||
- `dataset/agentic_data/tables/marketing.csv`
|
||||
- `dataset/agentic_data/tables/users.csv`
|
||||
- `dataset/agentic_data/train.jsonl`
|
||||
- `dataset/agentic_data/dev.jsonl`
|
||||
- `dataset/agentic_data/test.jsonl`
|
||||
- `dataset/agentic_data/manifest.json`
|
||||
|
||||
训练前校验：
|
||||
|
||||
```bash
|
||||
python scripts/validate_agentic_data.py \
|
||||
--data dataset/agentic_data/train.jsonl
|
||||
```
|
||||
|
||||
## SFT 冷启动
|
||||
|
||||
把专家轨迹转换成现有 `SFTDataset` 支持的 conversations 格式：
|
||||
|
||||
```bash
|
||||
python scripts/build_agentic_sft.py \
|
||||
--input dataset/agentic_data/train.jsonl \
|
||||
--output dataset/agentic_data/sft.jsonl
|
||||
```
|
||||
|
||||
使用已有 300M 预训练模型进行 Agent SFT：
|
||||
|
||||
```bash
|
||||
cd trainer
|
||||
|
||||
deepspeed --num_gpus=7 train_full_sft_deepspeed.py \
|
||||
--model_profile SearchLM-300M \
|
||||
--from_weight pretrain_searchlm \
|
||||
--save_weight agent_sft \
|
||||
--data_path ../dataset/agentic_data/sft.jsonl \
|
||||
--dtype float16 \
|
||||
--save_interval 200
|
||||
```
|
||||
|
||||
完整 14GB 通用 SFT 数据建议先训练一轮，再进入 Agent SFT。服务器可直接使用一键脚本：
|
||||
|
||||
```bash
|
||||
bash scripts/run_general_sft_300m.sh
|
||||
```
|
||||
|
||||
脚本默认使用 7 张 GPU、每卡 batch 1、梯度累积 8，有效 batch 为 56；如果检测到
|
||||
`checkpoints/general_sft_14g_1024_mla_latest`，会自动加 `--from_resume 1`。日志写入 `logs/`。
|
||||
|
||||
常用覆盖参数：
|
||||
|
||||
```bash
|
||||
USE_WANDB=1 SAVE_INTERVAL=200 bash scripts/run_general_sft_300m.sh
|
||||
```
|
||||
|
||||
通用 SFT 完成后，将下面 Agent SFT 命令中的 `--from_weight` 改为 `general_sft_14g`。
|
||||
|
||||
## Agentic RL
|
||||
|
||||
使用 trajectory-level reward 优化多步工具调用行为：
|
||||
|
||||
```bash
|
||||
cd trainer
|
||||
|
||||
deepspeed --num_gpus=7 train_agent_deepspeed.py \
|
||||
--model_profile SearchLM-300M \
|
||||
--from_weight agent_sft \
|
||||
--save_weight agent_grpo \
|
||||
--data_path ../dataset/agentic_data/train.jsonl \
|
||||
--num_generations 6 \
|
||||
--batch_size 1 \
|
||||
--accumulation_steps 4 \
|
||||
--dtype float16 \
|
||||
--save_interval 25
|
||||
```
|
||||
|
||||
断点续训：
|
||||
|
||||
```bash
|
||||
deepspeed --num_gpus=7 train_agent_deepspeed.py \
|
||||
--model_profile SearchLM-300M \
|
||||
--from_weight agent_sft \
|
||||
--save_weight agent_grpo \
|
||||
--data_path ../dataset/agentic_data/train.jsonl \
|
||||
--num_generations 6 \
|
||||
--batch_size 1 \
|
||||
--accumulation_steps 4 \
|
||||
--dtype float16 \
|
||||
--save_interval 25 \
|
||||
--from_resume 1
|
||||
```
|
||||
|
||||
`../checkpoints/` 保存完整 DeepSpeed 训练状态；`../out/agent_grpo_1024_mla.pth` 是轻量推理权重。
|
||||
|
||||
## Reward 指标
|
||||
|
||||
奖励函数位于 `agentic/data_analysis_env.py`，按整条轨迹评分：
|
||||
|
||||
- `task_success`：最终答案是否命中标准答案中的文本和数值校验。
|
||||
- `answer_score`：文本 contains 与数值 tolerance 的平均得分。
|
||||
- `tool_selection_f1`：调用工具集合与期望工具集合的 F1。
|
||||
- `schema_valid_rate`：工具名、JSON 和必填参数是否合法。
|
||||
- `exec_success_rate`：工具能否成功执行。
|
||||
- `invalid_action_rate`：非法 action 占比。
|
||||
- `avg_turns`：平均交互轮数。
|
||||
- `format_score`：`<tool_call>` 标签和 JSON 是否闭合合法。
|
||||
|
||||
## 离线评测
|
||||
|
||||
直接跑模型生成：
|
||||
|
||||
```bash
|
||||
python scripts/eval_agentic_data.py \
|
||||
--data dataset/agentic_data/test.jsonl \
|
||||
--weight agent_grpo \
|
||||
--save_dir out \
|
||||
--output reports/agentic_grpo_pred.jsonl \
|
||||
--json_output reports/agentic_grpo_metrics.json \
|
||||
--bad_cases_output reports/agentic_grpo_bad_cases.jsonl
|
||||
```
|
||||
|
||||
也可以评测已有预测文件：
|
||||
|
||||
```bash
|
||||
python scripts/eval_agentic_data.py \
|
||||
--data dataset/agentic_data/test.jsonl \
|
||||
--pred reports/agentic_grpo_pred.jsonl \
|
||||
--json_output reports/agentic_grpo_metrics.json
|
||||
```
|
||||
|
||||
建议固定对比：
|
||||
|
||||
```text
|
||||
pretrain_searchlm -> agent_sft -> agent_grpo
|
||||
```
|
||||
|
||||
核心简历指标使用固定 test split，不使用训练集结果。
|
||||
|
||||
## Demo
|
||||
|
||||
```bash
|
||||
streamlit run scripts/agentic_data_demo.py -- \
|
||||
--data dataset/agentic_data/dev.jsonl
|
||||
```
|
||||
|
||||
Demo 展示问题、数据表、指标文档、工具轨迹、最终答案和 reward 分解，适合面试时说明 Agent 行为如何被训练和评估。
|
||||
|
||||
## 简历表达
|
||||
|
||||
可以写成：
|
||||
|
||||
> 构建面向运营数据分析的 Agentic RL 系统，基于 300M Decoder-only 预训练模型完成工具轨迹 SFT 与 GRPO/CISPO 后训练；设计 retriever、python executor、calculator 交互环境和 trajectory-level reward，在固定测试集上评估任务成功率、工具选择准确率、无效 action 率、Python 执行成功率和平均交互轮数。
|
||||
|
|
@ -1,129 +1,69 @@
|
|||
{%- if tools %}
|
||||
{{- '<|im_start|>system
|
||||
' }}
|
||||
{{- '<|im_start|>system\n' }}
|
||||
{%- if messages[0].role == 'system' %}
|
||||
{{- messages[0].content + '
|
||||
|
||||
' }}
|
||||
{{- messages[0].content + '\n\n' }}
|
||||
{%- endif %}
|
||||
{{- "# Tools
|
||||
|
||||
You may call one or more functions to assist with the user query.
|
||||
|
||||
You are provided with function signatures within <tools></tools> XML tags:
|
||||
<tools>" }}
|
||||
{{- '# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>' }}
|
||||
{%- for tool in tools %}
|
||||
{{- "
|
||||
" }}
|
||||
{{- '\n' }}
|
||||
{{- tool | tojson }}
|
||||
{%- endfor %}
|
||||
{{- "
|
||||
</tools>
|
||||
|
||||
For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
|
||||
<tool_call>
|
||||
{"name": <function-name>, "arguments": <args-json-object>}
|
||||
</tool_call><|im_end|>
|
||||
" }}
|
||||
{{- '\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{"name": <function-name>, "arguments": <args-json-object>}\n</tool_call><|im_end|>\n' }}
|
||||
{%- else %}
|
||||
{%- if messages[0].role == 'system' %}
|
||||
{{- '<|im_start|>system
|
||||
' + messages[0].content + '<|im_end|>
|
||||
' }}
|
||||
{{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
|
||||
{%- for message in messages[::-1] %}
|
||||
{%- set index = (messages|length - 1) - loop.index0 %}
|
||||
{%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
|
||||
{%- set ns.multi_step_tool = false %}
|
||||
{%- set ns.last_query_index = index %}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- for message in messages %}
|
||||
{%- if message.content is string %}
|
||||
{%- set content = message.content %}
|
||||
{%- else %}
|
||||
{%- set content = '' %}
|
||||
{%- endif %}
|
||||
{%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
|
||||
{{- '<|im_start|>' + message.role + '
|
||||
' + content + '<|im_end|>' + '
|
||||
' }}
|
||||
{%- elif message.role == "assistant" %}
|
||||
{%- if message.role == 'user' or (message.role == 'system' and not loop.first) %}
|
||||
{{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>\n' }}
|
||||
{%- elif message.role == 'assistant' %}
|
||||
{%- set reasoning_content = '' %}
|
||||
{%- if message.reasoning_content is string %}
|
||||
{%- set reasoning_content = message.reasoning_content %}
|
||||
{%- else %}
|
||||
{%- if '</think>' in content %}
|
||||
{%- set reasoning_content = content.split('</think>')[0].rstrip('
|
||||
').split('<think>')[-1].lstrip('
|
||||
') %}
|
||||
{%- set content = content.split('</think>')[-1].lstrip('
|
||||
') %}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
{%- if true %}
|
||||
{{- '<|im_start|>' + message.role + '
|
||||
<think>
|
||||
' + reasoning_content.strip('
|
||||
') + '
|
||||
</think>
|
||||
|
||||
' + content.lstrip('
|
||||
') }}
|
||||
{%- elif '</think>' in content %}
|
||||
{%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
|
||||
{%- set content = content.split('</think>')[-1].lstrip('\n') %}
|
||||
{%- endif %}
|
||||
{{- '<|im_start|>assistant\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
|
||||
{%- if message.tool_calls %}
|
||||
{%- for tool_call in message.tool_calls %}
|
||||
{%- if (loop.first and content) or (not loop.first) %}
|
||||
{{- '
|
||||
' }}
|
||||
{%- if (loop.first and content) or not loop.first %}
|
||||
{{- '\n' }}
|
||||
{%- endif %}
|
||||
{%- if tool_call.function %}
|
||||
{%- set tool_call = tool_call.function %}
|
||||
{%- endif %}
|
||||
{{- '<tool_call>
|
||||
{"name": "' }}
|
||||
{{- tool_call.name }}
|
||||
{{- '", "arguments": ' }}
|
||||
{{- '<tool_call>\n{"name": "' + tool_call.name + '", "arguments": ' }}
|
||||
{%- if tool_call.arguments is string %}
|
||||
{{- tool_call.arguments }}
|
||||
{%- else %}
|
||||
{{- tool_call.arguments | tojson }}
|
||||
{%- endif %}
|
||||
{{- '}
|
||||
</tool_call>' }}
|
||||
{{- '}\n</tool_call>' }}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
{{- '<|im_end|>
|
||||
' }}
|
||||
{%- elif message.role == "tool" %}
|
||||
{%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
|
||||
{{- '<|im_end|>\n' }}
|
||||
{%- elif message.role == 'tool' %}
|
||||
{%- if loop.first or messages[loop.index0 - 1].role != 'tool' %}
|
||||
{{- '<|im_start|>user' }}
|
||||
{%- endif %}
|
||||
{{- '
|
||||
<tool_response>
|
||||
' }}
|
||||
{{- content }}
|
||||
{{- '
|
||||
</tool_response>' }}
|
||||
{%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
|
||||
{{- '<|im_end|>
|
||||
' }}
|
||||
{{- '\n<tool_response>\n' + content + '\n</tool_response>' }}
|
||||
{%- if loop.last or messages[loop.index0 + 1].role != 'tool' %}
|
||||
{{- '<|im_end|>\n' }}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- if add_generation_prompt %}
|
||||
{{- '<|im_start|>assistant
|
||||
' }}
|
||||
{{- '<|im_start|>assistant\n' }}
|
||||
{%- if open_thinking is defined and open_thinking is true %}
|
||||
{{- '<think>
|
||||
' }}
|
||||
{{- '<think>\n' }}
|
||||
{%- else %}
|
||||
{{- '<think>
|
||||
|
||||
</think>
|
||||
|
||||
' }}
|
||||
{{- '<think>\n\n</think>\n\n' }}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ ngrok==1.4.0
|
|||
nltk==3.8
|
||||
numpy==1.26.4
|
||||
openai==1.59.6
|
||||
pandas==2.2.3
|
||||
# peft==0.7.1
|
||||
psutil==5.9.8
|
||||
pydantic==2.11.5
|
||||
|
|
|
|||
|
|
@ -0,0 +1,97 @@
|
|||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pandas as pd
|
||||
import streamlit as st
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
|
||||
from agentic.data_analysis_env import extract_final_answer, score_agentic_trajectory
|
||||
from dataset.agentic_dataset import load_agentic_jsonl
|
||||
|
||||
|
||||
DEFAULT_DATA = "dataset/agentic_data/dev.jsonl"
|
||||
|
||||
|
||||
def trajectory_to_text(sample):
|
||||
chunks = []
|
||||
for item in sample.get("expert_trajectory", []):
|
||||
if item.get("role") == "assistant" and item.get("tool_call"):
|
||||
chunks.append("<tool_call>\n" + json.dumps(item["tool_call"], ensure_ascii=False, indent=2) + "\n</tool_call>")
|
||||
elif item.get("role") == "assistant":
|
||||
chunks.append(str(item.get("content", "")))
|
||||
return "\n---\n".join(chunks)
|
||||
|
||||
|
||||
def parse_turns(text):
|
||||
if "\n---\n" in text:
|
||||
return [part.strip() for part in text.split("\n---\n") if part.strip()]
|
||||
return [text.strip()] if text.strip() else []
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Agentic DataAnalysis Streamlit demo")
|
||||
parser.add_argument("--data", default=DEFAULT_DATA)
|
||||
parser.add_argument("--repo_root", default=".")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
st.set_page_config(page_title="Agentic DataAnalysis", layout="wide")
|
||||
st.title("Agentic DataAnalysis")
|
||||
|
||||
rows = load_agentic_jsonl(args.data) if os.path.exists(args.data) else []
|
||||
if not rows:
|
||||
st.error(f"数据不存在或为空：{args.data}")
|
||||
return
|
||||
|
||||
sample_ids = [row.get("id", str(i)) for i, row in enumerate(rows)]
|
||||
selected = st.sidebar.selectbox("样本", sample_ids)
|
||||
sample = rows[sample_ids.index(selected)]
|
||||
|
||||
st.subheader("问题")
|
||||
st.write(sample.get("question", ""))
|
||||
|
||||
left, right = st.columns([1.1, 1.0])
|
||||
with left:
|
||||
st.subheader("数据表")
|
||||
st.dataframe(pd.DataFrame(sample.get("tables", [])), use_container_width=True)
|
||||
st.subheader("指标文档")
|
||||
st.dataframe(pd.DataFrame(sample.get("documents", [])), use_container_width=True)
|
||||
|
||||
with right:
|
||||
st.subheader("标准答案")
|
||||
st.write(sample.get("answer", ""))
|
||||
st.json(sample.get("checks", {}))
|
||||
st.subheader("期望工具")
|
||||
st.write(", ".join(sample.get("expected_tools", [])))
|
||||
|
||||
st.subheader("轨迹与奖励")
|
||||
default_text = trajectory_to_text(sample)
|
||||
trajectory_text = st.text_area("轨迹文本（用 --- 分隔多个 assistant turn 可手动评测模型输出）", value=default_text, height=360)
|
||||
turn_outputs = [
|
||||
part
|
||||
for part in parse_turns(trajectory_text)
|
||||
if "<tool_response>" not in part
|
||||
]
|
||||
reward, parts = score_agentic_trajectory(turn_outputs, sample, repo_root=args.repo_root)
|
||||
|
||||
c1, c2, c3, c4, c5 = st.columns(5)
|
||||
c1.metric("Reward", f"{reward:.3f}")
|
||||
c2.metric("Success", f"{parts['task_success']:.3f}")
|
||||
c3.metric("Tool F1", f"{parts['tool_selection_f1']:.3f}")
|
||||
c4.metric("Invalid", f"{parts['invalid_action_rate']:.3f}")
|
||||
c5.metric("Turns", f"{parts['avg_turns']:.1f}")
|
||||
|
||||
st.subheader("Reward 分解")
|
||||
st.dataframe(pd.DataFrame([parts]).T.rename(columns={0: "score"}), use_container_width=True)
|
||||
|
||||
st.subheader("Final Answer")
|
||||
st.write(extract_final_answer(turn_outputs))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,488 @@
|
|||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import sys
|
||||
from collections import defaultdict
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
|
||||
from agentic.data_analysis_env import AgenticToolEnv, get_agentic_tools
|
||||
|
||||
|
||||
CHANNELS = ["抖音", "微信", "百度", "小红书", "自然流量"]
|
||||
CATEGORIES = ["美妆", "数码", "服饰", "食品", "家居"]
|
||||
PROVINCES = ["广东", "浙江", "北京", "上海", "四川", "湖北"]
|
||||
|
||||
METRIC_DOCS = [
|
||||
{
|
||||
"id": "metric_refund_amount",
|
||||
"title": "退款金额",
|
||||
"text": "退款金额使用 orders 表的 refund_amount 字段求和；按渠道、品类、省份或时间过滤后再聚合。",
|
||||
},
|
||||
{
|
||||
"id": "metric_net_revenue",
|
||||
"title": "净收入",
|
||||
"text": "净收入 = GMV - 退款金额。GMV 使用 orders.gmv，退款金额使用 orders.refund_amount。",
|
||||
},
|
||||
{
|
||||
"id": "metric_refund_rate",
|
||||
"title": "退款率",
|
||||
"text": "退款率 = 退款金额 / GMV。按品类或渠道对 refund_amount 与 gmv 分别求和后再相除。",
|
||||
},
|
||||
{
|
||||
"id": "metric_roas",
|
||||
"title": "ROAS",
|
||||
"text": "ROAS = GMV / 广告花费。GMV 来自 orders 表，广告花费来自 marketing 表的 spend 字段。",
|
||||
},
|
||||
{
|
||||
"id": "metric_new_users",
|
||||
"title": "新增用户",
|
||||
"text": "新增用户使用 users 表的 new_users 字段求和；活跃用户使用 active_users 字段求和。",
|
||||
},
|
||||
{
|
||||
"id": "schema_orders",
|
||||
"title": "orders 表字段",
|
||||
"text": "orders 包含 month、quarter、channel、category、province、gmv、refund_amount、orders_count。",
|
||||
},
|
||||
{
|
||||
"id": "schema_marketing",
|
||||
"title": "marketing 表字段",
|
||||
"text": "marketing 包含 month、quarter、channel、spend、impressions、clicks。",
|
||||
},
|
||||
{
|
||||
"id": "schema_users",
|
||||
"title": "users 表字段",
|
||||
"text": "users 包含 month、quarter、channel、province、active_users、new_users。",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def mkdir(path):
|
||||
os.makedirs(path, exist_ok=True)
|
||||
|
||||
|
||||
def write_csv(path, rows):
|
||||
if not rows:
|
||||
raise ValueError(f"empty csv rows: {path}")
|
||||
with open(path, "w", encoding="utf-8", newline="") as f:
|
||||
writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
|
||||
writer.writeheader()
|
||||
writer.writerows(rows)
|
||||
|
||||
|
||||
def relpath(path):
|
||||
cwd = os.getcwd()
|
||||
try:
|
||||
return os.path.relpath(path, cwd)
|
||||
except ValueError:
|
||||
return path
|
||||
|
||||
|
||||
def quarter_of(month):
|
||||
return f"Q{(month - 1) // 3 + 1}"
|
||||
|
||||
|
||||
def generate_tables(output_dir, seed):
|
||||
rng = random.Random(seed)
|
||||
table_dir = os.path.join(output_dir, "tables")
|
||||
mkdir(table_dir)
|
||||
|
||||
orders = []
|
||||
for month in range(1, 13):
|
||||
season = 1.0 + (0.18 if month in {6, 11, 12} else 0.0) + (0.08 if month in {1, 2} else 0.0)
|
||||
for channel_idx, channel in enumerate(CHANNELS):
|
||||
for category_idx, category in enumerate(CATEGORIES):
|
||||
for province_idx, province in enumerate(PROVINCES):
|
||||
base = 42000 + channel_idx * 4600 + category_idx * 3100 + province_idx * 1800
|
||||
gmv = round((base + rng.uniform(-5200, 6200)) * season, 2)
|
||||
refund_rate = 0.025 + 0.006 * category_idx + 0.003 * channel_idx + rng.uniform(0.0, 0.018)
|
||||
refund_amount = round(gmv * refund_rate, 2)
|
||||
orders_count = max(1, int(gmv / (90 + category_idx * 12 + rng.uniform(-8, 8))))
|
||||
orders.append(
|
||||
{
|
||||
"month": month,
|
||||
"quarter": quarter_of(month),
|
||||
"channel": channel,
|
||||
"category": category,
|
||||
"province": province,
|
||||
"gmv": gmv,
|
||||
"refund_amount": refund_amount,
|
||||
"orders_count": orders_count,
|
||||
}
|
||||
)
|
||||
|
||||
marketing = []
|
||||
for month in range(1, 13):
|
||||
for channel_idx, channel in enumerate(CHANNELS):
|
||||
spend = round(18000 + channel_idx * 3500 + month * 420 + rng.uniform(-1800, 2200), 2)
|
||||
clicks = max(1, int(spend / (2.2 + channel_idx * 0.25 + rng.uniform(-0.2, 0.25))))
|
||||
impressions = clicks * rng.randint(18, 36)
|
||||
marketing.append(
|
||||
{
|
||||
"month": month,
|
||||
"quarter": quarter_of(month),
|
||||
"channel": channel,
|
||||
"spend": spend,
|
||||
"impressions": impressions,
|
||||
"clicks": clicks,
|
||||
}
|
||||
)
|
||||
|
||||
users = []
|
||||
for month in range(1, 13):
|
||||
for channel_idx, channel in enumerate(CHANNELS):
|
||||
for province_idx, province in enumerate(PROVINCES):
|
||||
active_users = max(100, int(5200 + month * 110 + channel_idx * 520 + province_idx * 260 + rng.uniform(-420, 520)))
|
||||
new_users = max(20, int(active_users * (0.08 + channel_idx * 0.008 + rng.uniform(0.0, 0.04))))
|
||||
users.append(
|
||||
{
|
||||
"month": month,
|
||||
"quarter": quarter_of(month),
|
||||
"channel": channel,
|
||||
"province": province,
|
||||
"active_users": active_users,
|
||||
"new_users": new_users,
|
||||
}
|
||||
)
|
||||
|
||||
paths = {
|
||||
"orders": os.path.join(table_dir, "orders.csv"),
|
||||
"marketing": os.path.join(table_dir, "marketing.csv"),
|
||||
"users": os.path.join(table_dir, "users.csv"),
|
||||
}
|
||||
write_csv(paths["orders"], orders)
|
||||
write_csv(paths["marketing"], marketing)
|
||||
write_csv(paths["users"], users)
|
||||
return {"orders": orders, "marketing": marketing, "users": users, "paths": paths}
|
||||
|
||||
|
||||
def table_specs(paths):
|
||||
return [
|
||||
{
|
||||
"name": "orders",
|
||||
"path": relpath(paths["orders"]),
|
||||
"description": "订单聚合表，包含 GMV、退款金额和订单数。",
|
||||
},
|
||||
{
|
||||
"name": "marketing",
|
||||
"path": relpath(paths["marketing"]),
|
||||
"description": "投放聚合表，包含广告花费、曝光和点击。",
|
||||
},
|
||||
{
|
||||
"name": "users",
|
||||
"path": relpath(paths["users"]),
|
||||
"description": "用户活跃聚合表，包含活跃用户和新增用户。",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def sum_field(rows, field, **filters):
|
||||
total = 0.0
|
||||
for row in rows:
|
||||
if all(row.get(key) == value for key, value in filters.items()):
|
||||
total += float(row[field])
|
||||
return round(total, 2)
|
||||
|
||||
|
||||
def top_group(rows, group_key, value_field, filters):
|
||||
agg = defaultdict(float)
|
||||
for row in rows:
|
||||
if all(row.get(key) == value for key, value in filters.items()):
|
||||
agg[row[group_key]] += float(row[value_field])
|
||||
key, value = max(agg.items(), key=lambda item: item[1])
|
||||
return key, round(value, 2)
|
||||
|
||||
|
||||
def refund_rate_by_category(rows, month):
|
||||
agg = defaultdict(lambda: {"gmv": 0.0, "refund": 0.0})
|
||||
for row in rows:
|
||||
if row["month"] == month:
|
||||
agg[row["category"]]["gmv"] += float(row["gmv"])
|
||||
agg[row["category"]]["refund"] += float(row["refund_amount"])
|
||||
rates = {
|
||||
category: values["refund"] / values["gmv"]
|
||||
for category, values in agg.items()
|
||||
if values["gmv"] > 0
|
||||
}
|
||||
category, rate = max(rates.items(), key=lambda item: item[1])
|
||||
return category, round(rate, 4)
|
||||
|
||||
|
||||
def monthly_gmv_by_channel(orders, month, channel):
|
||||
return sum(float(row["gmv"]) for row in orders if row["month"] == month and row["channel"] == channel)
|
||||
|
||||
|
||||
def marketing_spend(marketing, month, channel):
|
||||
return sum(float(row["spend"]) for row in marketing if row["month"] == month and row["channel"] == channel)
|
||||
|
||||
|
||||
def build_tool_call(name, arguments):
|
||||
return {"name": name, "arguments": arguments}
|
||||
|
||||
|
||||
def make_sample(sample_id, question, answer, checks, tables, expected_tools, actions, task_type, max_turns=4):
|
||||
sample = {
|
||||
"id": sample_id,
|
||||
"question": question,
|
||||
"tables": tables,
|
||||
"documents": METRIC_DOCS,
|
||||
"tools": ["retriever", "python_executor", "calculator"],
|
||||
"expected_tools": expected_tools,
|
||||
"answer": answer,
|
||||
"checks": checks,
|
||||
"max_turns": max_turns,
|
||||
"metadata": {"task_type": task_type},
|
||||
}
|
||||
env = AgenticToolEnv(sample)
|
||||
trajectory = []
|
||||
for action in actions:
|
||||
trajectory.append({"role": "assistant", "tool_call": action})
|
||||
result = env.execute(action["name"], action["arguments"])
|
||||
trajectory.append({"role": "tool", "name": action["name"], "content": result})
|
||||
trajectory.append({"role": "assistant", "content": answer})
|
||||
sample["expert_trajectory"] = trajectory
|
||||
sample["tool_schemas"] = get_agentic_tools(sample["tools"])
|
||||
return sample
|
||||
|
||||
|
||||
def task_top_refund_channel(sample_id, tables, data, rng):
|
||||
quarter = rng.choice(["Q1", "Q2", "Q3", "Q4"])
|
||||
channel, value = top_group(data["orders"], "channel", "refund_amount", {"quarter": quarter})
|
||||
question = f"2024年{quarter}退款金额最高的渠道是什么？请给出退款金额。"
|
||||
code = (
|
||||
f'orders = pd.read_csv("{tables[0]["path"]}")\n'
|
||||
f'df = orders[orders["quarter"] == "{quarter}"]\n'
|
||||
'agg = df.groupby("channel")["refund_amount"].sum().sort_values(ascending=False)\n'
|
||||
'result = {"channel": agg.index[0], "refund_amount": round(float(agg.iloc[0]), 2)}'
|
||||
)
|
||||
answer = f"2024年{quarter}退款金额最高的渠道是{channel}，退款金额为{value:.2f}元。"
|
||||
return make_sample(
|
||||
sample_id,
|
||||
question,
|
||||
answer,
|
||||
{"contains": [channel], "number": value, "tolerance": 0.02},
|
||||
tables,
|
||||
["retriever", "python_executor"],
|
||||
[
|
||||
build_tool_call("retriever", {"query": "退款金额 指标口径 orders refund_amount", "top_k": 2}),
|
||||
build_tool_call("python_executor", {"code": code}),
|
||||
],
|
||||
"top_refund_channel",
|
||||
)
|
||||
|
||||
|
||||
def task_net_revenue(sample_id, tables, data, rng):
|
||||
month = rng.randint(1, 12)
|
||||
channel = rng.choice(CHANNELS)
|
||||
gmv = sum_field(data["orders"], "gmv", month=month, channel=channel)
|
||||
refund = sum_field(data["orders"], "refund_amount", month=month, channel=channel)
|
||||
value = round(gmv - refund, 2)
|
||||
question = f"2024年{month}月{channel}渠道的净收入是多少？"
|
||||
code = (
|
||||
f'orders = pd.read_csv("{tables[0]["path"]}")\n'
|
||||
f'df = orders[(orders["month"] == {month}) & (orders["channel"] == "{channel}")]\n'
|
||||
'result = {"net_revenue": round(float(df["gmv"].sum() - df["refund_amount"].sum()), 2)}'
|
||||
)
|
||||
answer = f"2024年{month}月{channel}渠道的净收入为{value:.2f}元。"
|
||||
return make_sample(
|
||||
sample_id,
|
||||
question,
|
||||
answer,
|
||||
{"contains": [channel], "number": value, "tolerance": 0.02},
|
||||
tables,
|
||||
["retriever", "python_executor"],
|
||||
[
|
||||
build_tool_call("retriever", {"query": "净收入 GMV 退款金额 计算口径", "top_k": 2}),
|
||||
build_tool_call("python_executor", {"code": code}),
|
||||
],
|
||||
"net_revenue",
|
||||
)
|
||||
|
||||
|
||||
def task_refund_rate_category(sample_id, tables, data, rng):
|
||||
month = rng.randint(1, 12)
|
||||
category, rate = refund_rate_by_category(data["orders"], month)
|
||||
percent = round(rate * 100, 2)
|
||||
question = f"2024年{month}月退款率最高的品类是什么？退款率是多少？"
|
||||
code = (
|
||||
f'orders = pd.read_csv("{tables[0]["path"]}")\n'
|
||||
f'df = orders[orders["month"] == {month}]\n'
|
||||
'agg = df.groupby("category")[["refund_amount", "gmv"]].sum()\n'
|
||||
'agg["refund_rate"] = agg["refund_amount"] / agg["gmv"]\n'
|
||||
'top = agg.sort_values("refund_rate", ascending=False).iloc[0]\n'
|
||||
'result = {"category": agg.sort_values("refund_rate", ascending=False).index[0], "refund_rate": round(float(top["refund_rate"] * 100), 2)}'
|
||||
)
|
||||
answer = f"2024年{month}月退款率最高的品类是{category}，退款率为{percent:.2f}%。"
|
||||
return make_sample(
|
||||
sample_id,
|
||||
question,
|
||||
answer,
|
||||
{"contains": [category], "number": percent, "tolerance": 0.02},
|
||||
tables,
|
||||
["retriever", "python_executor"],
|
||||
[
|
||||
build_tool_call("retriever", {"query": "退款率 指标口径 refund_amount gmv", "top_k": 2}),
|
||||
build_tool_call("python_executor", {"code": code}),
|
||||
],
|
||||
"refund_rate_category",
|
||||
)
|
||||
|
||||
|
||||
def task_roas(sample_id, tables, data, rng):
|
||||
month = rng.randint(1, 12)
|
||||
channel = rng.choice([c for c in CHANNELS if c != "自然流量"])
|
||||
gmv = monthly_gmv_by_channel(data["orders"], month, channel)
|
||||
spend = marketing_spend(data["marketing"], month, channel)
|
||||
value = round(gmv / spend, 2)
|
||||
question = f"2024年{month}月{channel}渠道的ROAS是多少？"
|
||||
code = (
|
||||
f'orders = pd.read_csv("{tables[0]["path"]}")\n'
|
||||
f'marketing = pd.read_csv("{tables[1]["path"]}")\n'
|
||||
f'gmv = orders[(orders["month"] == {month}) & (orders["channel"] == "{channel}")]["gmv"].sum()\n'
|
||||
f'spend = marketing[(marketing["month"] == {month}) & (marketing["channel"] == "{channel}")]["spend"].sum()\n'
|
||||
'result = {"roas": round(float(gmv / spend), 2)}'
|
||||
)
|
||||
answer = f"2024年{month}月{channel}渠道的ROAS为{value:.2f}。"
|
||||
return make_sample(
|
||||
sample_id,
|
||||
question,
|
||||
answer,
|
||||
{"contains": [channel], "number": value, "tolerance": 0.02},
|
||||
tables,
|
||||
["retriever", "python_executor"],
|
||||
[
|
||||
build_tool_call("retriever", {"query": "ROAS GMV 广告花费 指标口径", "top_k": 2}),
|
||||
build_tool_call("python_executor", {"code": code}),
|
||||
],
|
||||
"roas",
|
||||
)
|
||||
|
||||
|
||||
def task_new_users(sample_id, tables, data, rng):
|
||||
month = rng.randint(1, 12)
|
||||
channel = rng.choice(CHANNELS)
|
||||
province, value = top_group(data["users"], "province", "new_users", {"month": month, "channel": channel})
|
||||
question = f"2024年{month}月{channel}渠道新增用户最多的省份是哪个？新增用户数是多少？"
|
||||
code = (
|
||||
f'users = pd.read_csv("{tables[2]["path"]}")\n'
|
||||
f'df = users[(users["month"] == {month}) & (users["channel"] == "{channel}")]\n'
|
||||
'agg = df.groupby("province")["new_users"].sum().sort_values(ascending=False)\n'
|
||||
'result = {"province": agg.index[0], "new_users": int(agg.iloc[0])}'
|
||||
)
|
||||
answer = f"2024年{month}月{channel}渠道新增用户最多的省份是{province}，新增用户数为{int(value)}。"
|
||||
return make_sample(
|
||||
sample_id,
|
||||
question,
|
||||
answer,
|
||||
{"contains": [province], "number": int(value), "tolerance": 0.01},
|
||||
tables,
|
||||
["retriever", "python_executor"],
|
||||
[
|
||||
build_tool_call("retriever", {"query": "新增用户 users new_users 字段口径", "top_k": 2}),
|
||||
build_tool_call("python_executor", {"code": code}),
|
||||
],
|
||||
"new_users",
|
||||
)
|
||||
|
||||
|
||||
def task_incremental_gmv(sample_id, tables, data, rng):
|
||||
month = rng.randint(1, 12)
|
||||
channel = rng.choice([c for c in CHANNELS if c != "自然流量"])
|
||||
pct = rng.choice([10, 12, 15, 20])
|
||||
gmv = monthly_gmv_by_channel(data["orders"], month, channel)
|
||||
spend = marketing_spend(data["marketing"], month, channel)
|
||||
roas = gmv / spend
|
||||
incremental = round(spend * pct / 100 * roas, 2)
|
||||
question = f"如果2024年{month}月{channel}渠道广告花费增加{pct}%，且ROAS保持不变，预计GMV增加多少？"
|
||||
code = (
|
||||
f'orders = pd.read_csv("{tables[0]["path"]}")\n'
|
||||
f'marketing = pd.read_csv("{tables[1]["path"]}")\n'
|
||||
f'gmv = orders[(orders["month"] == {month}) & (orders["channel"] == "{channel}")]["gmv"].sum()\n'
|
||||
f'spend = marketing[(marketing["month"] == {month}) & (marketing["channel"] == "{channel}")]["spend"].sum()\n'
|
||||
'result = {"spend": round(float(spend), 2), "roas": round(float(gmv / spend), 4)}'
|
||||
)
|
||||
expression = f"{spend:.2f} * {pct} / 100 * {roas:.6f}"
|
||||
answer = f"预计GMV增加{incremental:.2f}元。"
|
||||
return make_sample(
|
||||
sample_id,
|
||||
question,
|
||||
answer,
|
||||
{"contains": ["GMV"], "number": incremental, "tolerance": 0.05},
|
||||
tables,
|
||||
["retriever", "python_executor", "calculator"],
|
||||
[
|
||||
build_tool_call("retriever", {"query": "ROAS GMV 广告花费 预算增加 计算口径", "top_k": 2}),
|
||||
build_tool_call("python_executor", {"code": code}),
|
||||
build_tool_call("calculator", {"expression": expression}),
|
||||
],
|
||||
"incremental_gmv",
|
||||
max_turns=5,
|
||||
)
|
||||
|
||||
|
||||
TASK_BUILDERS = [
|
||||
task_top_refund_channel,
|
||||
task_net_revenue,
|
||||
task_refund_rate_category,
|
||||
task_roas,
|
||||
task_new_users,
|
||||
task_incremental_gmv,
|
||||
]
|
||||
|
||||
|
||||
def write_jsonl(path, rows):
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
for row in rows:
|
||||
f.write(json.dumps(row, ensure_ascii=False) + "\n")
|
||||
|
||||
|
||||
def build_split(name, size, tables, data, seed):
|
||||
rng = random.Random(seed)
|
||||
rows = []
|
||||
for i in range(size):
|
||||
task_fn = TASK_BUILDERS[i % len(TASK_BUILDERS)]
|
||||
sample_id = f"agentic_{name}_{i:06d}"
|
||||
rows.append(task_fn(sample_id, tables, data, rng))
|
||||
rng.shuffle(rows)
|
||||
return rows
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Build Agentic DataAnalysis tasks")
|
||||
parser.add_argument("--output_dir", default="dataset/agentic_data")
|
||||
parser.add_argument("--train_size", type=int, default=30000)
|
||||
parser.add_argument("--dev_size", type=int, default=1000)
|
||||
parser.add_argument("--test_size", type=int, default=1000)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
mkdir(args.output_dir)
|
||||
data = generate_tables(args.output_dir, args.seed)
|
||||
tables = table_specs(data["paths"])
|
||||
splits = {
|
||||
"train": build_split("train", args.train_size, tables, data, args.seed + 1),
|
||||
"dev": build_split("dev", args.dev_size, tables, data, args.seed + 2),
|
||||
"test": build_split("test", args.test_size, tables, data, args.seed + 3),
|
||||
}
|
||||
for split, rows in splits.items():
|
||||
write_jsonl(os.path.join(args.output_dir, f"{split}.jsonl"), rows)
|
||||
manifest = {
|
||||
"task": "AgenticDataAnalysis",
|
||||
"seed": args.seed,
|
||||
"splits": {split: len(rows) for split, rows in splits.items()},
|
||||
"tables": tables,
|
||||
"tools": [tool["function"]["name"] for tool in get_agentic_tools()],
|
||||
"task_types": [fn.__name__.replace("task_", "") for fn in TASK_BUILDERS],
|
||||
}
|
||||
with open(os.path.join(args.output_dir, "manifest.json"), "w", encoding="utf-8") as f:
|
||||
json.dump(manifest, f, ensure_ascii=False, indent=2)
|
||||
print(json.dumps(manifest, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,86 @@
|
|||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
|
||||
from agentic.data_analysis_env import AGENTIC_SYSTEM_PROMPT, format_agentic_user_prompt, get_agentic_tools
|
||||
from dataset.agentic_dataset import load_agentic_jsonl
|
||||
|
||||
|
||||
def assistant_tool_message(tool_call):
|
||||
return {
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"reasoning_content": "",
|
||||
"tools": "",
|
||||
"tool_calls": json.dumps(
|
||||
[
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tool_call.get("name", ""),
|
||||
"arguments": tool_call.get("arguments", {}),
|
||||
},
|
||||
}
|
||||
],
|
||||
ensure_ascii=False,
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def normal_message(role, content):
|
||||
return {
|
||||
"role": role,
|
||||
"content": content,
|
||||
"reasoning_content": "",
|
||||
"tools": "",
|
||||
"tool_calls": "",
|
||||
}
|
||||
|
||||
|
||||
def convert_sample(sample):
|
||||
tools = get_agentic_tools(sample.get("tools"))
|
||||
conversations = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": AGENTIC_SYSTEM_PROMPT,
|
||||
"reasoning_content": "",
|
||||
"tools": json.dumps(tools, ensure_ascii=False),
|
||||
"tool_calls": "",
|
||||
},
|
||||
normal_message("user", format_agentic_user_prompt(sample)),
|
||||
]
|
||||
for item in sample.get("expert_trajectory", []):
|
||||
if item.get("role") == "assistant" and item.get("tool_call"):
|
||||
conversations.append(assistant_tool_message(item["tool_call"]))
|
||||
elif item.get("role") == "tool":
|
||||
conversations.append(normal_message("tool", json.dumps(item.get("content", {}), ensure_ascii=False)))
|
||||
elif item.get("role") == "assistant":
|
||||
conversations.append(normal_message("assistant", str(item.get("content", ""))))
|
||||
return {"id": sample.get("id", ""), "conversations": conversations}
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Build Agentic DataAnalysis SFT conversations")
|
||||
parser.add_argument("--input", required=True)
|
||||
parser.add_argument("--output", required=True)
|
||||
parser.add_argument("--limit", type=int, default=0)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
rows = load_agentic_jsonl(args.input)
|
||||
if args.limit > 0:
|
||||
rows = rows[:args.limit]
|
||||
os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
|
||||
with open(args.output, "w", encoding="utf-8") as f:
|
||||
for sample in rows:
|
||||
f.write(json.dumps(convert_sample(sample), ensure_ascii=False) + "\n")
|
||||
print(json.dumps({"input": args.input, "output": args.output, "samples": len(rows)}, ensure_ascii=False))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,237 @@
|
|||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
|
||||
from agentic.data_analysis_env import (
|
||||
AgenticToolEnv,
|
||||
average_agentic_metrics,
|
||||
extract_final_answer,
|
||||
format_agentic_user_prompt,
|
||||
get_agentic_tools,
|
||||
parse_tool_calls,
|
||||
score_agentic_trajectory,
|
||||
)
|
||||
from dataset.agentic_dataset import load_agentic_jsonl
|
||||
|
||||
|
||||
def apply_eval_model_profile(args):
|
||||
if not args.model_profile:
|
||||
return
|
||||
path = args.model_profiles_path
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
profiles = json.load(f)
|
||||
if args.model_profile not in profiles:
|
||||
raise ValueError(f"Unknown model_profile={args.model_profile}. Available: {', '.join(profiles)}")
|
||||
profile = profiles[args.model_profile]
|
||||
for key, value in profile.items():
|
||||
if hasattr(args, key):
|
||||
setattr(args, key, value)
|
||||
|
||||
|
||||
def load_model(args):
|
||||
import torch
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from model.model_minimind import MiniMindForCausalLM
|
||||
from model.model_minimind_mla import MiniMindMLAForCausalLM
|
||||
from trainer.trainer_utils import build_lm_config, get_model_suffix, resolve_attention_type
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
|
||||
config = build_lm_config(
|
||||
hidden_size=args.hidden_size,
|
||||
num_hidden_layers=args.num_hidden_layers,
|
||||
max_seq_len=args.max_total_len,
|
||||
use_moe=bool(args.use_moe),
|
||||
attention_type=resolve_attention_type(args),
|
||||
kv_lora_rank=args.kv_lora_rank,
|
||||
q_lora_rank=args.q_lora_rank,
|
||||
rope_dim=args.rope_dim,
|
||||
num_attention_heads=args.num_attention_heads,
|
||||
num_key_value_heads=args.num_key_value_heads,
|
||||
intermediate_size=args.intermediate_size,
|
||||
)
|
||||
model = MiniMindMLAForCausalLM(config) if args.attention_type == "mla" or args.use_mla else MiniMindForCausalLM(config)
|
||||
suffix = get_model_suffix(config)
|
||||
weight_path = os.path.join(args.save_dir, f"{args.weight}_{config.hidden_size}{suffix}.pth")
|
||||
weights = torch.load(weight_path, map_location=args.device)
|
||||
model.load_state_dict(weights, strict=False)
|
||||
dtype = torch.float16 if args.dtype == "float16" else torch.float32
|
||||
if args.dtype == "bfloat16":
|
||||
dtype = torch.bfloat16
|
||||
model = model.to(args.device)
|
||||
if args.device != "cpu" and args.dtype != "float32":
|
||||
model = model.to(dtype=dtype)
|
||||
return model.eval(), tokenizer
|
||||
|
||||
|
||||
def generate_once(model, tokenizer, messages, tools, args):
|
||||
import torch
|
||||
|
||||
prompt = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True,
|
||||
tools=tools,
|
||||
open_thinking=bool(args.open_thinking),
|
||||
)
|
||||
inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(args.device)
|
||||
if inputs["input_ids"].size(1) > args.max_total_len:
|
||||
inputs["input_ids"] = inputs["input_ids"][:, -args.max_total_len:]
|
||||
inputs["attention_mask"] = inputs["attention_mask"][:, -args.max_total_len:]
|
||||
st = time.time()
|
||||
with torch.no_grad():
|
||||
generated = model.generate(
|
||||
inputs["input_ids"],
|
||||
attention_mask=inputs["attention_mask"],
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
do_sample=bool(args.do_sample),
|
||||
temperature=args.temperature,
|
||||
top_p=args.top_p,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
eos_token_id=tokenizer.eos_token_id,
|
||||
)
|
||||
text = tokenizer.decode(generated[0][inputs["input_ids"].size(1):], skip_special_tokens=True)
|
||||
return text, time.time() - st
|
||||
|
||||
|
||||
def run_episode(model, tokenizer, sample, args):
|
||||
tools = get_agentic_tools(sample.get("tools"))
|
||||
messages = [
|
||||
{"role": "system", "content": "你是一个面向运营数据分析的 Agent。"},
|
||||
{"role": "user", "content": format_agentic_user_prompt(sample)},
|
||||
]
|
||||
env = AgenticToolEnv(sample, repo_root=args.repo_root, timeout=args.tool_timeout)
|
||||
turn_outputs = []
|
||||
elapsed = 0.0
|
||||
unfinished = False
|
||||
for turn in range(args.max_turns):
|
||||
text, dt = generate_once(model, tokenizer, messages, tools, args)
|
||||
elapsed += dt
|
||||
turn_outputs.append(text)
|
||||
calls = parse_tool_calls(text)
|
||||
if not calls:
|
||||
break
|
||||
unfinished = turn == args.max_turns - 1
|
||||
messages.append({"role": "assistant", "content": text})
|
||||
for call in calls:
|
||||
result = env.execute(call.get("name", ""), call.get("arguments", {}))
|
||||
messages.append({"role": "tool", "content": json.dumps(result, ensure_ascii=False)})
|
||||
if unfinished:
|
||||
break
|
||||
reward, parts = score_agentic_trajectory(
|
||||
turn_outputs,
|
||||
sample,
|
||||
repo_root=args.repo_root,
|
||||
unfinished=unfinished,
|
||||
execute_tools=True,
|
||||
)
|
||||
return {
|
||||
"id": sample.get("id", ""),
|
||||
"question": sample.get("question", ""),
|
||||
"prediction": extract_final_answer(turn_outputs),
|
||||
"turn_outputs": turn_outputs,
|
||||
"reward": reward,
|
||||
"metrics": parts,
|
||||
"latency_sec": elapsed,
|
||||
}
|
||||
|
||||
|
||||
def eval_predictions(args, rows):
|
||||
pred_map = {}
|
||||
with open(args.pred, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
if line.strip():
|
||||
item = json.loads(line)
|
||||
pred_map[item.get("id")] = item
|
||||
outputs = []
|
||||
for sample in rows:
|
||||
item = pred_map.get(sample.get("id"), {})
|
||||
turns = item.get("turn_outputs") or [item.get("prediction") or item.get("response") or item.get("output") or ""]
|
||||
reward, parts = score_agentic_trajectory(turns, sample, repo_root=args.repo_root)
|
||||
outputs.append(
|
||||
{
|
||||
"id": sample.get("id", ""),
|
||||
"question": sample.get("question", ""),
|
||||
"prediction": extract_final_answer(turns),
|
||||
"turn_outputs": turns,
|
||||
"reward": reward,
|
||||
"metrics": parts,
|
||||
}
|
||||
)
|
||||
return outputs
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Evaluate Agentic DataAnalysis model or predictions")
|
||||
parser.add_argument("--data", required=True)
|
||||
parser.add_argument("--pred", default="")
|
||||
parser.add_argument("--output", default="reports/agentic_eval_predictions.jsonl")
|
||||
parser.add_argument("--json_output", default="reports/agentic_eval_metrics.json")
|
||||
parser.add_argument("--bad_cases_output", default="reports/agentic_eval_bad_cases.jsonl")
|
||||
parser.add_argument("--max_samples", type=int, default=0)
|
||||
parser.add_argument("--repo_root", default=".")
|
||||
parser.add_argument("--device", default="cuda")
|
||||
parser.add_argument("--save_dir", default="out")
|
||||
parser.add_argument("--weight", default="agent_grpo")
|
||||
parser.add_argument("--tokenizer_path", default="model")
|
||||
parser.add_argument("--dtype", default="float16", choices=["float16", "bfloat16", "float32"])
|
||||
parser.add_argument("--hidden_size", default=1024, type=int)
|
||||
parser.add_argument("--num_hidden_layers", default=24, type=int)
|
||||
parser.add_argument("--num_attention_heads", default=8, type=int)
|
||||
parser.add_argument("--num_key_value_heads", default=4, type=int)
|
||||
parser.add_argument("--intermediate_size", default=3264, type=int)
|
||||
parser.add_argument("--use_moe", default=0, type=int, choices=[0, 1])
|
||||
parser.add_argument("--attention_type", default="mla", choices=["gqa", "mha", "mqa", "mla"])
|
||||
parser.add_argument("--use_mla", default=0, type=int, choices=[0, 1])
|
||||
parser.add_argument("--kv_lora_rank", default=192, type=int)
|
||||
parser.add_argument("--q_lora_rank", default=384, type=int)
|
||||
parser.add_argument("--rope_dim", default=None, type=int)
|
||||
parser.add_argument("--max_total_len", default=2500, type=int)
|
||||
parser.add_argument("--max_new_tokens", default=512, type=int)
|
||||
parser.add_argument("--max_turns", default=4, type=int)
|
||||
parser.add_argument("--tool_timeout", default=3, type=int)
|
||||
parser.add_argument("--temperature", default=0.8, type=float)
|
||||
parser.add_argument("--top_p", default=0.9, type=float)
|
||||
parser.add_argument("--do_sample", default=1, type=int)
|
||||
parser.add_argument("--open_thinking", default=0, type=int)
|
||||
parser.add_argument("--model_profile", type=str, default="")
|
||||
parser.add_argument("--model_profiles_path", type=str, default="configs/searchlm_profiles.json")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
apply_eval_model_profile(args)
|
||||
rows = load_agentic_jsonl(args.data)
|
||||
if args.max_samples > 0:
|
||||
rows = rows[: args.max_samples]
|
||||
|
||||
if args.pred:
|
||||
outputs = eval_predictions(args, rows)
|
||||
else:
|
||||
model, tokenizer = load_model(args)
|
||||
outputs = [run_episode(model, tokenizer, sample, args) for sample in rows]
|
||||
|
||||
os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
|
||||
os.makedirs(os.path.dirname(args.json_output) or ".", exist_ok=True)
|
||||
os.makedirs(os.path.dirname(args.bad_cases_output) or ".", exist_ok=True)
|
||||
with open(args.output, "w", encoding="utf-8") as f:
|
||||
for item in outputs:
|
||||
f.write(json.dumps(item, ensure_ascii=False) + "\n")
|
||||
metrics = average_agentic_metrics([item["metrics"] for item in outputs])
|
||||
metrics["samples"] = len(outputs)
|
||||
with open(args.json_output, "w", encoding="utf-8") as f:
|
||||
json.dump(metrics, f, ensure_ascii=False, indent=2)
|
||||
with open(args.bad_cases_output, "w", encoding="utf-8") as f:
|
||||
for item in outputs:
|
||||
if item["metrics"].get("task_success", 0) < 1:
|
||||
f.write(json.dumps(item, ensure_ascii=False) + "\n")
|
||||
print(json.dumps(metrics, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,78 @@
|
|||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
|
||||
DATA_PATH="${DATA_PATH:-${REPO_ROOT}/dataset/sft_t2t.jsonl}"
|
||||
WEIGHT_PATH="${WEIGHT_PATH:-${REPO_ROOT}/out/pretrain_searchlm_1024_mla.pth}"
|
||||
SAVE_WEIGHT="${SAVE_WEIGHT:-general_sft_14g}"
|
||||
NUM_GPUS="${NUM_GPUS:-7}"
|
||||
EPOCHS="${EPOCHS:-1}"
|
||||
MICRO_BATCH="${MICRO_BATCH:-1}"
|
||||
ACCUMULATION_STEPS="${ACCUMULATION_STEPS:-8}"
|
||||
LEARNING_RATE="${LEARNING_RATE:-1e-5}"
|
||||
SAVE_INTERVAL="${SAVE_INTERVAL:-500}"
|
||||
LOG_INTERVAL="${LOG_INTERVAL:-20}"
|
||||
NUM_WORKERS="${NUM_WORKERS:-8}"
|
||||
USE_WANDB="${USE_WANDB:-0}"
|
||||
FORCE_RESUME="${FORCE_RESUME:-auto}"
|
||||
|
||||
if [[ ! -s "${DATA_PATH}" ]]; then
|
||||
echo "SFT dataset not found or empty: ${DATA_PATH}" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [[ ! -s "${WEIGHT_PATH}" ]]; then
|
||||
echo "Pretrained checkpoint not found: ${WEIGHT_PATH}" >&2
|
||||
echo "Expected from_weight=pretrain_searchlm -> out/pretrain_searchlm_1024_mla.pth" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if ! command -v deepspeed >/dev/null 2>&1; then
|
||||
echo "deepspeed command not found. Install it in the server training environment first." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
mkdir -p "${REPO_ROOT}/logs" "${REPO_ROOT}/out" "${REPO_ROOT}/checkpoints"
|
||||
|
||||
LATEST_FILE="${REPO_ROOT}/checkpoints/${SAVE_WEIGHT}_1024_mla_latest"
|
||||
RESUME=0
|
||||
if [[ "${FORCE_RESUME}" == "1" ]]; then
|
||||
RESUME=1
|
||||
elif [[ "${FORCE_RESUME}" == "auto" && -s "${LATEST_FILE}" ]]; then
|
||||
RESUME=1
|
||||
fi
|
||||
|
||||
STAMP="$(date +%Y%m%d_%H%M%S)"
|
||||
LOG_FILE="${REPO_ROOT}/logs/${SAVE_WEIGHT}_${STAMP}.log"
|
||||
|
||||
echo "repo_root=${REPO_ROOT}"
|
||||
echo "data_path=${DATA_PATH}"
|
||||
echo "weight_path=${WEIGHT_PATH}"
|
||||
echo "save_weight=${SAVE_WEIGHT}"
|
||||
echo "num_gpus=${NUM_GPUS}"
|
||||
echo "effective_batch=$((NUM_GPUS * MICRO_BATCH * ACCUMULATION_STEPS))"
|
||||
echo "resume=${RESUME}"
|
||||
echo "log_file=${LOG_FILE}"
|
||||
|
||||
WANDB_ARGS=()
|
||||
if [[ "${USE_WANDB}" == "1" ]]; then
|
||||
WANDB_ARGS+=(--use_wandb)
|
||||
fi
|
||||
|
||||
cd "${REPO_ROOT}/trainer"
|
||||
|
||||
deepspeed --num_gpus="${NUM_GPUS}" train_full_sft_deepspeed.py \
|
||||
--model_profile SearchLM-300M \
|
||||
--from_weight pretrain_searchlm \
|
||||
--save_weight "${SAVE_WEIGHT}" \
|
||||
--data_path "${DATA_PATH}" \
|
||||
--epochs "${EPOCHS}" \
|
||||
--batch_size "${MICRO_BATCH}" \
|
||||
--accumulation_steps "${ACCUMULATION_STEPS}" \
|
||||
--learning_rate "${LEARNING_RATE}" \
|
||||
--dtype float16 \
|
||||
--num_workers "${NUM_WORKERS}" \
|
||||
--log_interval "${LOG_INTERVAL}" \
|
||||
--save_interval "${SAVE_INTERVAL}" \
|
||||
--from_resume "${RESUME}" \
|
||||
"${WANDB_ARGS[@]}" 2>&1 | tee "${LOG_FILE}"
|
||||
|
|
@ -0,0 +1,69 @@
|
|||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
|
||||
from agentic.data_analysis_env import score_agentic_trajectory
|
||||
from dataset.agentic_dataset import load_agentic_jsonl
|
||||
|
||||
|
||||
REQUIRED_FIELDS = ["id", "question", "tables", "documents", "tools", "expected_tools", "answer", "checks"]
|
||||
|
||||
|
||||
def trajectory_to_turn_outputs(sample):
|
||||
turns = []
|
||||
for item in sample.get("expert_trajectory", []):
|
||||
if item.get("role") == "assistant" and item.get("tool_call"):
|
||||
turns.append("<tool_call>\n" + json.dumps(item["tool_call"], ensure_ascii=False) + "\n</tool_call>")
|
||||
elif item.get("role") == "assistant":
|
||||
turns.append(str(item.get("content", "")))
|
||||
return turns
|
||||
|
||||
|
||||
def validate_sample(sample):
|
||||
errors = []
|
||||
for field in REQUIRED_FIELDS:
|
||||
if field not in sample:
|
||||
errors.append(f"missing field: {field}")
|
||||
for table in sample.get("tables", []):
|
||||
path = table.get("path", "")
|
||||
if not path:
|
||||
errors.append("table path is empty")
|
||||
elif not os.path.exists(path):
|
||||
errors.append(f"table path not found: {path}")
|
||||
if not sample.get("expert_trajectory"):
|
||||
errors.append("missing expert_trajectory")
|
||||
if not errors:
|
||||
reward, parts = score_agentic_trajectory(trajectory_to_turn_outputs(sample), sample)
|
||||
if parts.get("task_success", 0) < 1:
|
||||
errors.append(f"expert trajectory not successful: reward={reward:.3f}, parts={parts}")
|
||||
return errors
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Validate Agentic DataAnalysis JSONL")
|
||||
parser.add_argument("--data", required=True)
|
||||
parser.add_argument("--max_errors", type=int, default=20)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
rows = load_agentic_jsonl(args.data)
|
||||
total_errors = 0
|
||||
for row in rows:
|
||||
errors = validate_sample(row)
|
||||
if errors:
|
||||
total_errors += len(errors)
|
||||
if total_errors <= args.max_errors:
|
||||
print(json.dumps({"id": row.get("id"), "errors": errors}, ensure_ascii=False))
|
||||
summary = {"data": args.data, "samples": len(rows), "errors": total_errors}
|
||||
print(json.dumps(summary, ensure_ascii=False))
|
||||
if total_errors:
|
||||
raise SystemExit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,576 @@
|
|||
"""
|
||||
DeepSpeed Agentic DataAnalysis GRPO/CISPO 入口。
|
||||
|
||||
推荐在 trainer/ 目录启动：
|
||||
deepspeed --num_gpus=7 train_agent_deepspeed.py --model_profile SearchLM-300M
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
|
||||
__package__ = "trainer"
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import random
|
||||
import warnings
|
||||
from contextlib import nullcontext
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.nn.functional as F
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
|
||||
from agentic.data_analysis_env import average_agentic_metrics, parse_tool_calls, score_agentic_trajectory
|
||||
from dataset.agentic_dataset import AgenticRLDataset
|
||||
from trainer.rollout_engine import compute_per_token_logps, create_rollout_engine
|
||||
from trainer.trainer_utils import (
|
||||
Logger,
|
||||
SkipBatchSampler,
|
||||
add_deepspeed_args,
|
||||
add_model_profile_args,
|
||||
apply_model_profile,
|
||||
assert_no_training_output_collision,
|
||||
build_lm_config,
|
||||
get_cuda_peak_memory_gb,
|
||||
init_model,
|
||||
is_main_process,
|
||||
load_deepspeed_checkpoint,
|
||||
load_deepspeed_config,
|
||||
log_training_setup,
|
||||
reduce_metrics,
|
||||
resolve_attention_type,
|
||||
save_deepspeed_checkpoint,
|
||||
save_model_weights,
|
||||
setup_seed,
|
||||
)
|
||||
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
|
||||
def init_deepspeed_runtime(args, deepspeed):
|
||||
deepspeed.init_distributed(dist_backend="nccl")
|
||||
local_rank = int(os.environ.get("LOCAL_RANK", args.local_rank if args.local_rank >= 0 else 0))
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.set_device(local_rank)
|
||||
args.device = f"cuda:{local_rank}"
|
||||
else:
|
||||
args.device = "cpu"
|
||||
args.training_backend = "DeepSpeed-ZeRO"
|
||||
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else local_rank))
|
||||
return local_rank
|
||||
|
||||
|
||||
def collate_fn(batch):
|
||||
return {
|
||||
"id": [x["id"] for x in batch],
|
||||
"messages": [x["messages"] for x in batch],
|
||||
"tools": [x["tools"] for x in batch],
|
||||
"sample": [x["sample"] for x in batch],
|
||||
}
|
||||
|
||||
|
||||
def rollout_single(args, rollout_engine, tokenizer, messages, tools, sample):
|
||||
all_outputs = []
|
||||
prompt_ids = None
|
||||
response_ids = []
|
||||
response_mask = []
|
||||
response_old_logps = []
|
||||
final_context = ""
|
||||
unfinished = False
|
||||
open_thinking = random.random() < args.thinking_ratio
|
||||
|
||||
for turn in range(args.max_turns):
|
||||
context = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True,
|
||||
tools=tools,
|
||||
open_thinking=open_thinking,
|
||||
)
|
||||
inputs = tokenizer(context, return_tensors="pt", add_special_tokens=False).to(args.device)
|
||||
context_ids = inputs["input_ids"][0].tolist()
|
||||
if prompt_ids is None:
|
||||
prompt_ids = context_ids
|
||||
|
||||
rollout_result = rollout_engine.rollout(
|
||||
prompt_ids=inputs["input_ids"],
|
||||
attention_mask=inputs["attention_mask"],
|
||||
num_generations=1,
|
||||
max_new_tokens=args.max_gen_len,
|
||||
temperature=args.temperature,
|
||||
)
|
||||
new_ids = rollout_result.completion_ids[0].tolist()
|
||||
new_logps = rollout_result.per_token_logps[0].tolist()
|
||||
pairs = [
|
||||
(token_id, logp)
|
||||
for token_id, logp in zip(new_ids, new_logps)
|
||||
if token_id != tokenizer.pad_token_id and token_id != tokenizer.eos_token_id
|
||||
]
|
||||
new_ids = [token_id for token_id, _ in pairs]
|
||||
new_logps = [logp for _, logp in pairs]
|
||||
new_text = rollout_result.completions[0]
|
||||
all_outputs.append(new_text)
|
||||
response_ids.extend(new_ids)
|
||||
response_mask.extend([1] * len(new_ids))
|
||||
response_old_logps.extend(new_logps)
|
||||
final_context = context + new_text
|
||||
|
||||
calls = parse_tool_calls(new_text)
|
||||
if not calls:
|
||||
break
|
||||
|
||||
unfinished = turn == args.max_turns - 1
|
||||
messages.append({"role": "assistant", "content": new_text})
|
||||
from agentic.data_analysis_env import AgenticToolEnv
|
||||
|
||||
env = AgenticToolEnv(sample, repo_root=args.repo_root, timeout=args.tool_timeout)
|
||||
for call in calls:
|
||||
result = env.execute(call.get("name", ""), call.get("arguments", {}))
|
||||
result_str = json.dumps(result, ensure_ascii=False)[: args.max_observation_chars]
|
||||
messages.append({"role": "tool", "content": result_str})
|
||||
|
||||
observe_context = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=not unfinished,
|
||||
tools=tools,
|
||||
open_thinking=open_thinking,
|
||||
)
|
||||
observe_ids = tokenizer(observe_context, return_tensors="pt", add_special_tokens=False)["input_ids"][0].tolist()
|
||||
current_len = len(prompt_ids) + len(response_ids)
|
||||
obs_delta = observe_ids[current_len:]
|
||||
response_ids.extend(obs_delta)
|
||||
response_mask.extend([0] * len(obs_delta))
|
||||
response_old_logps.extend([0.0] * len(obs_delta))
|
||||
final_context = observe_context
|
||||
if unfinished:
|
||||
break
|
||||
|
||||
return (
|
||||
final_context,
|
||||
prompt_ids or [],
|
||||
response_ids,
|
||||
response_mask,
|
||||
response_old_logps,
|
||||
all_outputs,
|
||||
unfinished,
|
||||
)
|
||||
|
||||
|
||||
def rollout_batch(args, rollout_engine, tokenizer, messages_batch, tools_batch, samples):
|
||||
contexts = []
|
||||
prompt_ids_batch = []
|
||||
response_ids_batch = []
|
||||
response_masks_batch = []
|
||||
response_old_logps_batch = []
|
||||
turn_outputs_batch = []
|
||||
unfinished_batch = []
|
||||
for messages, tools, sample in zip(messages_batch, tools_batch, samples):
|
||||
for _ in range(args.num_generations):
|
||||
messages_copy = [dict(message) for message in messages]
|
||||
context, prompt_ids, response_ids, response_mask, old_logps, turn_outputs, unfinished = rollout_single(
|
||||
args,
|
||||
rollout_engine,
|
||||
tokenizer,
|
||||
messages_copy,
|
||||
tools,
|
||||
sample,
|
||||
)
|
||||
contexts.append(context)
|
||||
prompt_ids_batch.append(prompt_ids)
|
||||
response_ids_batch.append(response_ids)
|
||||
response_masks_batch.append(response_mask)
|
||||
response_old_logps_batch.append(old_logps)
|
||||
turn_outputs_batch.append(turn_outputs)
|
||||
unfinished_batch.append(unfinished)
|
||||
return contexts, prompt_ids_batch, response_ids_batch, response_masks_batch, response_old_logps_batch, turn_outputs_batch, unfinished_batch
|
||||
|
||||
|
||||
def calculate_agentic_rewards(args, turn_outputs_batch, samples, unfinished_batch):
|
||||
rewards = []
|
||||
metrics = []
|
||||
for idx, turn_outputs in enumerate(turn_outputs_batch):
|
||||
sample_idx = idx // args.num_generations
|
||||
reward, parts = score_agentic_trajectory(
|
||||
turn_outputs,
|
||||
samples[sample_idx],
|
||||
repo_root=args.repo_root,
|
||||
unfinished=unfinished_batch[idx],
|
||||
execute_tools=True,
|
||||
)
|
||||
rewards.append(reward)
|
||||
metrics.append(parts)
|
||||
return torch.tensor(rewards, dtype=torch.float32, device=args.device), average_agentic_metrics(metrics)
|
||||
|
||||
|
||||
def pack_rollout_samples(args, tokenizer, prompt_ids_batch, response_ids_batch, response_masks_batch, response_old_logps_batch):
|
||||
packed = []
|
||||
for prompt_ids, response_ids, response_mask, old_logps in zip(
|
||||
prompt_ids_batch,
|
||||
response_ids_batch,
|
||||
response_masks_batch,
|
||||
response_old_logps_batch,
|
||||
):
|
||||
ids = prompt_ids + response_ids
|
||||
mask = [0] * len(prompt_ids) + response_mask
|
||||
old_per_token = [0.0] * max(len(prompt_ids) - 1, 0) + old_logps
|
||||
if len(ids) > args.max_total_len:
|
||||
ids = ids[-args.max_total_len:]
|
||||
mask = mask[-args.max_total_len:]
|
||||
old_per_token = old_per_token[-(len(ids) - 1):]
|
||||
target_len = max(len(ids) - 1, 0)
|
||||
if len(old_per_token) < target_len:
|
||||
old_per_token += [0.0] * (target_len - len(old_per_token))
|
||||
elif len(old_per_token) > target_len:
|
||||
old_per_token = old_per_token[-target_len:]
|
||||
packed.append((ids, mask, old_per_token))
|
||||
|
||||
max_len = max(len(ids) for ids, _, _ in packed)
|
||||
input_ids = torch.tensor(
|
||||
[ids + [tokenizer.pad_token_id] * (max_len - len(ids)) for ids, _, _ in packed],
|
||||
device=args.device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
full_response_masks = torch.tensor(
|
||||
[mask + [0] * (max_len - len(mask)) for _, mask, _ in packed],
|
||||
device=args.device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
old_per_token_logps = torch.tensor(
|
||||
[old + [0.0] * ((max_len - 1) - len(old)) for _, _, old in packed],
|
||||
device=args.device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
return input_ids, full_response_masks, old_per_token_logps
|
||||
|
||||
|
||||
def train_epoch(args, lm_config, model_engine, tokenizer, loader, iters, rollout_engine, ref_model, start_step=0, wandb=None):
|
||||
for step, batch in enumerate(loader, start=start_step + 1):
|
||||
contexts, prompt_ids_batch, response_ids_batch, response_masks_batch, response_old_logps_batch, turn_outputs_batch, unfinished_batch = rollout_batch(
|
||||
args,
|
||||
rollout_engine,
|
||||
tokenizer,
|
||||
batch["messages"],
|
||||
batch["tools"],
|
||||
batch["sample"],
|
||||
)
|
||||
input_ids, full_response_masks, old_per_token_logps = pack_rollout_samples(
|
||||
args,
|
||||
tokenizer,
|
||||
prompt_ids_batch,
|
||||
response_ids_batch,
|
||||
response_masks_batch,
|
||||
response_old_logps_batch,
|
||||
)
|
||||
full_mask = (input_ids != tokenizer.pad_token_id).long()
|
||||
rewards, reward_stats = calculate_agentic_rewards(args, turn_outputs_batch, batch["sample"], unfinished_batch)
|
||||
|
||||
res = model_engine(input_ids, attention_mask=full_mask)
|
||||
aux_loss = res.aux_loss if getattr(lm_config, "use_moe", False) and res.aux_loss is not None else torch.tensor(0.0, device=args.device)
|
||||
per_token_logps = F.log_softmax(res.logits[:, :-1, :], dim=-1).gather(2, input_ids[:, 1:].unsqueeze(-1)).squeeze(-1)
|
||||
|
||||
with torch.no_grad():
|
||||
ref_per_token_logps = compute_per_token_logps(
|
||||
ref_model,
|
||||
input_ids,
|
||||
input_ids.size(1) - 1,
|
||||
attention_mask=full_mask,
|
||||
)
|
||||
|
||||
completion_mask = full_response_masks[:, 1:]
|
||||
token_counts = completion_mask.sum(dim=1)
|
||||
valid_rows = token_counts > 0
|
||||
|
||||
grouped_rewards = rewards.view(-1, args.num_generations)
|
||||
mean_r = grouped_rewards.mean(dim=1).repeat_interleave(args.num_generations)
|
||||
std_r = grouped_rewards.std(dim=1, unbiased=False).repeat_interleave(args.num_generations)
|
||||
advantages = (rewards - mean_r) / (std_r + 1e-4)
|
||||
|
||||
kl_div = ref_per_token_logps - per_token_logps
|
||||
per_token_kl = torch.exp(kl_div) - kl_div - 1
|
||||
ratio = torch.exp(per_token_logps - old_per_token_logps)
|
||||
if args.loss_type == "cispo":
|
||||
clamped_ratio = torch.clamp(ratio, max=args.epsilon_high).detach()
|
||||
per_token_loss = -(clamped_ratio * advantages.unsqueeze(1) * per_token_logps - args.beta * per_token_kl)
|
||||
else:
|
||||
clipped_ratio = torch.clamp(ratio, 1 - args.epsilon, 1 + args.epsilon)
|
||||
per_token_loss1 = ratio * advantages.unsqueeze(1)
|
||||
per_token_loss2 = clipped_ratio * advantages.unsqueeze(1)
|
||||
per_token_loss = -(torch.min(per_token_loss1, per_token_loss2) - args.beta * per_token_kl)
|
||||
policy_loss = (
|
||||
((per_token_loss * completion_mask).sum(dim=1)[valid_rows] / token_counts[valid_rows].clamp(min=1)).mean()
|
||||
if valid_rows.any()
|
||||
else per_token_loss.sum() * 0.0
|
||||
)
|
||||
loss = policy_loss + aux_loss
|
||||
|
||||
model_engine.backward(loss)
|
||||
model_engine.step()
|
||||
|
||||
if step % args.log_interval == 0 or step == iters:
|
||||
reward_val = rewards.mean().item()
|
||||
avg_len_val = token_counts.float().mean().item()
|
||||
kl_ref_val = ((ref_per_token_logps - per_token_logps) * completion_mask).sum().item() / max(token_counts.sum().item(), 1)
|
||||
lr_list = model_engine.get_lr()
|
||||
current_lr = lr_list[0] if lr_list else args.learning_rate
|
||||
metrics = reduce_metrics(
|
||||
{
|
||||
"reward": reward_val,
|
||||
"avg_len": avg_len_val,
|
||||
"kl_ref": kl_ref_val,
|
||||
"policy_loss": policy_loss.item(),
|
||||
"aux_loss": aux_loss.item(),
|
||||
**reward_stats,
|
||||
}
|
||||
)
|
||||
peak_memory_gb = get_cuda_peak_memory_gb()
|
||||
memory_text = f", peak_memory: {peak_memory_gb:.2f} GB" if peak_memory_gb is not None else ""
|
||||
Logger(
|
||||
f"Epoch:[{args.current_epoch + 1}/{args.epochs}]({step}/{iters}), "
|
||||
f"Reward:{metrics['reward']:.4f}, Success:{metrics.get('task_success', 0):.4f}, "
|
||||
f"ToolF1:{metrics.get('tool_selection_f1', 0):.4f}, Invalid:{metrics.get('invalid_action_rate', 0):.4f}, "
|
||||
f"Exec:{metrics.get('exec_success_rate', 0):.4f}, Turns:{metrics.get('avg_turns', 0):.2f}, "
|
||||
f"KL:{metrics['kl_ref']:.4f}, Loss:{metrics['policy_loss']:.4f}, LR:{current_lr:.8f}"
|
||||
f"{memory_text}"
|
||||
)
|
||||
if wandb and is_main_process():
|
||||
log_data = {
|
||||
"reward": metrics["reward"],
|
||||
"kl_ref": metrics["kl_ref"],
|
||||
"policy_loss": metrics["policy_loss"],
|
||||
"aux_loss": metrics["aux_loss"],
|
||||
"avg_response_len": metrics["avg_len"],
|
||||
"learning_rate": current_lr,
|
||||
}
|
||||
log_data.update({key: value for key, value in metrics.items() if key not in log_data})
|
||||
wandb.log(log_data)
|
||||
|
||||
if args.debug_mode and is_main_process() and step % args.debug_interval == 0:
|
||||
for i, sample_id in enumerate(batch["id"][:2]):
|
||||
Logger(f"[DEBUG] sample={sample_id}")
|
||||
for j in range(args.num_generations):
|
||||
idx = i * args.num_generations + j
|
||||
Logger(f"[DEBUG] gen={j}, reward={rewards[idx].item():.4f}, unfinished={unfinished_batch[idx]}")
|
||||
Logger("\n---\n".join(turn_outputs_batch[idx]))
|
||||
Logger("=" * 80)
|
||||
|
||||
if step % args.save_interval == 0 or step == iters:
|
||||
model_engine.eval()
|
||||
save_deepspeed_checkpoint(
|
||||
model_engine,
|
||||
lm_config,
|
||||
weight=args.save_weight,
|
||||
epoch=args.current_epoch,
|
||||
step=step,
|
||||
wandb=wandb,
|
||||
save_dir="../checkpoints",
|
||||
)
|
||||
if is_main_process():
|
||||
ckp = save_model_weights(lm_config, model_engine, save_dir=args.save_dir, weight=args.save_weight)
|
||||
Logger(f"Lightweight model saved: {ckp}")
|
||||
model_engine.train()
|
||||
rollout_engine.update_policy(model_engine)
|
||||
|
||||
del contexts, prompt_ids_batch, response_ids_batch, response_masks_batch, response_old_logps_batch
|
||||
del turn_outputs_batch, unfinished_batch, input_ids, full_response_masks, old_per_token_logps
|
||||
del per_token_logps, ref_per_token_logps, rewards, grouped_rewards, mean_r, std_r, advantages
|
||||
del completion_mask, token_counts, res, loss, policy_loss, aux_loss
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Agentic DataAnalysis GRPO/CISPO with DeepSpeed ZeRO")
|
||||
parser.add_argument("--save_dir", type=str, default="../out")
|
||||
parser.add_argument("--save_weight", default="agent_grpo", type=str)
|
||||
parser.add_argument("--epochs", type=int, default=1)
|
||||
parser.add_argument("--batch_size", type=int, default=1)
|
||||
parser.add_argument("--learning_rate", type=float, default=3e-7)
|
||||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
|
||||
parser.add_argument("--dtype", type=str, default="float16", choices=["float16", "bfloat16", "float32"])
|
||||
parser.add_argument("--num_workers", type=int, default=4)
|
||||
parser.add_argument("--accumulation_steps", type=int, default=4)
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0)
|
||||
parser.add_argument("--log_interval", type=int, default=1)
|
||||
parser.add_argument("--save_interval", type=int, default=25)
|
||||
parser.add_argument("--hidden_size", default=768, type=int)
|
||||
parser.add_argument("--num_hidden_layers", default=8, type=int)
|
||||
parser.add_argument("--num_attention_heads", default=8, type=int)
|
||||
parser.add_argument("--num_key_value_heads", default=4, type=int)
|
||||
parser.add_argument("--intermediate_size", default=None, type=int)
|
||||
parser.add_argument("--use_moe", default=0, type=int, choices=[0, 1])
|
||||
parser.add_argument("--attention_type", default="gqa", choices=["gqa", "mha", "mqa", "mla"])
|
||||
parser.add_argument("--use_mla", default=0, type=int, choices=[0, 1])
|
||||
parser.add_argument("--kv_lora_rank", default=128, type=int)
|
||||
parser.add_argument("--q_lora_rank", default=256, type=int)
|
||||
parser.add_argument("--rope_dim", default=None, type=int)
|
||||
parser.add_argument("--max_seq_len", default=1536, type=int)
|
||||
parser.add_argument("--max_gen_len", type=int, default=512)
|
||||
parser.add_argument("--max_total_len", type=int, default=2500)
|
||||
parser.add_argument("--max_turns", type=int, default=4)
|
||||
parser.add_argument("--max_observation_chars", type=int, default=2048)
|
||||
parser.add_argument("--tool_timeout", type=int, default=3)
|
||||
parser.add_argument("--repo_root", type=str, default="..")
|
||||
parser.add_argument("--data_path", type=str, default="../dataset/agentic_data/train.jsonl")
|
||||
parser.add_argument("--num_generations", type=int, default=6)
|
||||
parser.add_argument("--beta", type=float, default=0.05)
|
||||
parser.add_argument("--loss_type", type=str, default="cispo", choices=["grpo", "cispo"])
|
||||
parser.add_argument("--epsilon", type=float, default=0.2)
|
||||
parser.add_argument("--epsilon_high", type=float, default=5.0)
|
||||
parser.add_argument("--temperature", type=float, default=0.8)
|
||||
parser.add_argument("--thinking_ratio", type=float, default=0.1)
|
||||
parser.add_argument("--from_weight", default="agent_sft", type=str)
|
||||
parser.add_argument("--from_resume", default=0, type=int, choices=[0, 1])
|
||||
parser.add_argument("--overwrite_output", default=0, type=int, choices=[0, 1])
|
||||
parser.add_argument("--use_wandb", action="store_true")
|
||||
parser.add_argument("--wandb_project", type=str, default="Agentic-DataAnalysis-RL-DS")
|
||||
parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1])
|
||||
parser.add_argument("--debug_mode", action="store_true")
|
||||
parser.add_argument("--debug_interval", type=int, default=20)
|
||||
parser.add_argument("--rollout_engine", type=str, default="torch", choices=["torch", "sglang"])
|
||||
parser.add_argument("--sglang_base_url", type=str, default="http://localhost:8998")
|
||||
parser.add_argument("--sglang_model_path", type=str, default="../model")
|
||||
parser.add_argument("--sglang_shared_path", type=str, default="./sglang_ckpt_agentic")
|
||||
add_model_profile_args(parser)
|
||||
add_deepspeed_args(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
model_profile = apply_model_profile(args)
|
||||
|
||||
try:
|
||||
import deepspeed
|
||||
except ImportError as exc:
|
||||
raise RuntimeError("DeepSpeed未安装。服务器环境执行: pip install deepspeed") from exc
|
||||
|
||||
local_rank = init_deepspeed_runtime(args, deepspeed)
|
||||
os.makedirs(args.save_dir, exist_ok=True)
|
||||
|
||||
lm_config = build_lm_config(
|
||||
hidden_size=args.hidden_size,
|
||||
num_hidden_layers=args.num_hidden_layers,
|
||||
max_seq_len=max(args.max_seq_len + args.max_gen_len, args.max_total_len),
|
||||
use_moe=bool(args.use_moe),
|
||||
attention_type=resolve_attention_type(args),
|
||||
kv_lora_rank=args.kv_lora_rank,
|
||||
q_lora_rank=args.q_lora_rank,
|
||||
rope_dim=args.rope_dim,
|
||||
num_attention_heads=args.num_attention_heads,
|
||||
num_key_value_heads=args.num_key_value_heads,
|
||||
intermediate_size=args.intermediate_size,
|
||||
)
|
||||
assert_no_training_output_collision(
|
||||
lm_config,
|
||||
weight=args.save_weight,
|
||||
save_dir=args.save_dir,
|
||||
checkpoint_dir="../checkpoints",
|
||||
from_resume=args.from_resume,
|
||||
overwrite=args.overwrite_output,
|
||||
)
|
||||
|
||||
model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
|
||||
if args.use_compile == 1:
|
||||
model = torch.compile(model)
|
||||
Logger("torch.compile enabled before DeepSpeed initialize")
|
||||
|
||||
train_ds = AgenticRLDataset(args.data_path, tokenizer, max_length=args.max_total_len)
|
||||
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
|
||||
loader_for_count = DataLoader(train_ds, batch_size=args.batch_size, sampler=train_sampler, collate_fn=collate_fn)
|
||||
iters = len(loader_for_count)
|
||||
total_optimizer_steps = max(1, math.ceil(iters / args.accumulation_steps) * args.epochs)
|
||||
ds_config = load_deepspeed_config(args, total_optimizer_steps)
|
||||
|
||||
model_engine, _, _, _ = deepspeed.initialize(
|
||||
model=model,
|
||||
model_parameters=model.parameters(),
|
||||
config=ds_config,
|
||||
)
|
||||
|
||||
ckp_data = None
|
||||
if args.from_resume == 1:
|
||||
ckp_data = load_deepspeed_checkpoint(
|
||||
model_engine,
|
||||
lm_config,
|
||||
weight=args.save_weight,
|
||||
save_dir="../checkpoints",
|
||||
)
|
||||
|
||||
ref_model, _ = init_model(lm_config, args.from_weight, device=args.device)
|
||||
ref_model = ref_model.eval().requires_grad_(False)
|
||||
if torch.cuda.is_available() and args.dtype in {"float16", "bfloat16"}:
|
||||
ref_dtype = torch.float16 if args.dtype == "float16" else torch.bfloat16
|
||||
ref_model = ref_model.to(dtype=ref_dtype)
|
||||
|
||||
device_type = "cuda" if "cuda" in args.device else "cpu"
|
||||
amp_dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
|
||||
autocast_ctx = nullcontext() if device_type == "cpu" or args.dtype == "float32" else torch.cuda.amp.autocast(dtype=amp_dtype)
|
||||
|
||||
rollout_engine = create_rollout_engine(
|
||||
engine_type=args.rollout_engine,
|
||||
policy_model=model_engine,
|
||||
tokenizer=tokenizer,
|
||||
device=args.device,
|
||||
autocast_ctx=autocast_ctx,
|
||||
sglang_base_url=args.sglang_base_url,
|
||||
sglang_model_path=args.sglang_model_path,
|
||||
sglang_shared_path=args.sglang_shared_path,
|
||||
)
|
||||
rollout_engine.update_policy(model_engine)
|
||||
|
||||
wandb = None
|
||||
if args.use_wandb and is_main_process():
|
||||
import swanlab as wandb
|
||||
|
||||
wandb_id = ckp_data.get("wandb_id") if ckp_data else None
|
||||
wandb.init(
|
||||
project=args.wandb_project,
|
||||
name=f"Agentic-RL-DS-{args.attention_type}-{args.hidden_size}",
|
||||
id=wandb_id,
|
||||
resume="must" if wandb_id else None,
|
||||
)
|
||||
|
||||
log_training_setup(
|
||||
args,
|
||||
lm_config,
|
||||
stage="agentic_rl_deepspeed",
|
||||
dataset_len=len(train_ds),
|
||||
iters=iters,
|
||||
tokens_per_sample=args.max_total_len,
|
||||
extra={
|
||||
"model_profile": args.model_profile if model_profile else "",
|
||||
"deepspeed_config": args.deepspeed_config,
|
||||
"zero_stage": ds_config.get("zero_optimization", {}).get("stage", "none"),
|
||||
"num_generations": args.num_generations,
|
||||
"rollout_engine": args.rollout_engine,
|
||||
"loss_type": args.loss_type,
|
||||
"beta": args.beta,
|
||||
"max_turns": args.max_turns,
|
||||
"total_optimizer_steps": total_optimizer_steps,
|
||||
},
|
||||
)
|
||||
|
||||
start_epoch = ckp_data.get("epoch", 0) if ckp_data else 0
|
||||
start_step = ckp_data.get("step", 0) if ckp_data else 0
|
||||
|
||||
for epoch in range(start_epoch, args.epochs):
|
||||
args.current_epoch = epoch
|
||||
train_sampler and train_sampler.set_epoch(epoch)
|
||||
setup_seed(42 + epoch + local_rank)
|
||||
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_fn,
|
||||
)
|
||||
if skip > 0:
|
||||
Logger(f"Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始")
|
||||
train_epoch(args, lm_config, model_engine, tokenizer, loader, len(loader) + skip, rollout_engine, ref_model, skip, wandb)
|
||||
|
||||
if dist.is_initialized():
|
||||
dist.destroy_process_group()
|
||||
|
|
@ -71,6 +71,29 @@ def validate_sft_supervision(train_ds, max_samples=32):
|
|||
)
|
||||
|
||||
|
||||
def validate_chat_template(tokenizer):
|
||||
try:
|
||||
rendered = tokenizer.apply_chat_template(
|
||||
[
|
||||
{"role": "user", "content": "template check"},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
],
|
||||
tokenize=False,
|
||||
add_generation_prompt=False,
|
||||
)
|
||||
except Exception as exc:
|
||||
raise RuntimeError(
|
||||
"Tokenizer chat template is invalid. Check model/chat_template.jinja and "
|
||||
"model/tokenizer_config.json before loading the SFT dataset."
|
||||
) from exc
|
||||
required_markers = ("<|im_start|>user", "<|im_start|>assistant", "<|im_end|>")
|
||||
missing = [marker for marker in required_markers if marker not in rendered]
|
||||
if missing:
|
||||
raise RuntimeError(f"Tokenizer chat template is missing required markers: {missing}")
|
||||
if is_main_process():
|
||||
Logger("Tokenizer chat template validation passed")
|
||||
|
||||
|
||||
def init_deepspeed_runtime(args, deepspeed):
|
||||
deepspeed.init_distributed(dist_backend="nccl")
|
||||
local_rank = int(os.environ.get("LOCAL_RANK", args.local_rank if args.local_rank >= 0 else 0))
|
||||
|
|
@ -129,7 +152,12 @@ def train_epoch(args, lm_config, model_engine, loader, iters, start_step=0, wand
|
|||
}
|
||||
)
|
||||
|
||||
if step % args.save_interval == 0 or step == iters:
|
||||
should_save = step % args.save_interval == 0 or step == iters
|
||||
del input_ids, labels, res, loss, aux_loss
|
||||
|
||||
if should_save:
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
model_engine.eval()
|
||||
save_deepspeed_checkpoint(
|
||||
model_engine,
|
||||
|
|
@ -144,10 +172,13 @@ def train_epoch(args, lm_config, model_engine, loader, iters, start_step=0, wand
|
|||
if is_main_process():
|
||||
ckp = save_model_weights(lm_config, model_engine, save_dir=args.save_dir, weight=args.save_weight)
|
||||
Logger(f"Lightweight model saved: {ckp}")
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
# rank0 导出轻量权重时，其他 rank 必须等待，避免提前销毁 NCCL process group。
|
||||
if dist.is_initialized():
|
||||
dist.barrier()
|
||||
model_engine.train()
|
||||
|
||||
del input_ids, labels, res, loss, aux_loss
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Full SFT with DeepSpeed ZeRO")
|
||||
|
|
@ -222,6 +253,7 @@ if __name__ == "__main__":
|
|||
)
|
||||
|
||||
model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
|
||||
validate_chat_template(tokenizer)
|
||||
if args.use_compile == 1:
|
||||
model = torch.compile(model)
|
||||
Logger("torch.compile enabled before DeepSpeed initialize")
|
||||
|
|
|
|||
Loading…
Reference in New Issue