Add training and planning eval exports

2026-05-11 14:50:00 +08:00
parent 8cca3793e0
commit be731caed7
15 changed files with 1074 additions and 4 deletions
@@ -19,6 +19,8 @@ from typing import Any

 DEFAULT_REQUEST_ID = "aabbccdd"
 DEFAULT_TIMESTAMP_STEP_MS = 60_000
+DEFAULT_SYSTEM_PROMPT = "你是小爱同学，中文智能语音助手。"
+DEFAULT_CONTEXT_FIELDS = ("location", "rag")


 class ProductDataError(ValueError):
@@ -205,7 +207,7 @@ def case_to_record(
    if not current_query:
        raise ProductDataError(f"case {index} current 用户 line must be non-empty")

-    target = str(case.get("target") or "").strip()
+    target = canonical_target(str(case.get("target") or "").strip())
    if not target:
        raise ProductDataError(f"case {index} target is required")

@@ -447,6 +449,89 @@ def export_dataset_table(
    }


+def export_training_jsonl(
+    records: list[dict[str, Any]],
+    *,
+    output_path: str,
+    session_num: int = 5,
+    session_time_minutes: int = 5,
+    context_fields: list[str] | None = None,
+    system_prompt: str = DEFAULT_SYSTEM_PROMPT,
+    require_validation_ok: bool = True,
+    overwrite: bool = True,
+) -> dict[str, Any]:
+    if session_num <= 0:
+        raise ProductDataError("session_num must be greater than 0")
+    if session_time_minutes <= 0:
+        raise ProductDataError("session_time_minutes must be greater than 0")
+    validation = validate_dataset_records(records)
+    if require_validation_ok and not validation["ok"]:
+        raise ProductDataError(f"records failed validation with {validation['error_count']} errors")
+    path = Path(output_path).expanduser()
+    if path.exists() and not overwrite:
+        raise ProductDataError(f"output_path already exists: {output_path}")
+    path.parent.mkdir(parents=True, exist_ok=True)
+    fields = normalize_context_fields(context_fields)
+    content = "".join(
+        training_jsonl_line(
+            record,
+            session_num=session_num,
+            session_time_minutes=session_time_minutes,
+            context_fields=fields,
+            system_prompt=system_prompt,
+        )
+        + "\n"
+        for record in records
+    )
+    path.write_text(content, encoding="utf-8")
+    return {
+        "output_path": str(path),
+        "output_format": "jsonl",
+        "record_count": len(records),
+        "bytes_written": len(content.encode("utf-8")),
+        "validation": validation,
+    }
+
+
+def export_planning_eval_csv(
+    records: list[dict[str, Any]],
+    *,
+    output_path: str,
+    session_num: int = 5,
+    session_time_minutes: int = 5,
+    context_fields: list[str] | None = None,
+    complex_default: bool = False,
+    require_validation_ok: bool = True,
+    overwrite: bool = True,
+) -> dict[str, Any]:
+    if session_num <= 0:
+        raise ProductDataError("session_num must be greater than 0")
+    if session_time_minutes <= 0:
+        raise ProductDataError("session_time_minutes must be greater than 0")
+    validation = validate_dataset_records(records)
+    if require_validation_ok and not validation["ok"]:
+        raise ProductDataError(f"records failed validation with {validation['error_count']} errors")
+    path = Path(output_path).expanduser()
+    if path.exists() and not overwrite:
+        raise ProductDataError(f"output_path already exists: {output_path}")
+    path.parent.mkdir(parents=True, exist_ok=True)
+    content = render_planning_eval_csv(
+        records,
+        session_num=session_num,
+        session_time_minutes=session_time_minutes,
+        context_fields=context_fields,
+        complex_default=complex_default,
+    )
+    path.write_text(content, encoding="utf-8-sig")
+    return {
+        "output_path": str(path),
+        "output_format": "csv",
+        "record_count": len(records),
+        "bytes_written": len(content.encode("utf-8-sig")),
+        "validation": validation,
+    }
+
+
 def render_dataset_records_table_csv(records: list[dict[str, Any]]) -> str:
    output = StringIO()
    writer = csv.DictWriter(output, fieldnames=dataset_table_columns(), lineterminator="\n")
@@ -456,6 +541,185 @@ def render_dataset_records_table_csv(records: list[dict[str, Any]]) -> str:
    return output.getvalue()


+def render_planning_eval_csv(
+    records: list[dict[str, Any]],
+    *,
+    session_num: int = 5,
+    session_time_minutes: int = 5,
+    context_fields: list[str] | None = None,
+    complex_default: bool = False,
+) -> str:
+    fields = normalize_context_fields(context_fields)
+    output = StringIO()
+    writer = csv.DictWriter(
+        output,
+        fieldnames=["request_id", "newPrompt", "query", "类别真实标签", "code标签", "complex"],
+        lineterminator="\n",
+    )
+    writer.writeheader()
+    for record in records:
+        source = record.get("source") if isinstance(record.get("source"), dict) else {}
+        turn = record.get("turn") if isinstance(record.get("turn"), dict) else {}
+        label = record.get("label") if isinstance(record.get("label"), dict) else {}
+        target = str(label.get("target") or "")
+        writer.writerow(
+            {
+                "request_id": str(source.get("request_id") or ""),
+                "newPrompt": build_training_instruction(
+                    record,
+                    session_num=session_num,
+                    session_time_minutes=session_time_minutes,
+                    context_fields=fields,
+                ),
+                "query": str(turn.get("query") or ""),
+                "类别真实标签": category_label_from_target(target),
+                "code标签": target,
+                "complex": "true" if complex_default else "false",
+            }
+        )
+    return output.getvalue()
+
+
+def build_training_instruction(
+    record: dict[str, Any],
+    *,
+    session_num: int = 5,
+    session_time_minutes: int = 5,
+    context_fields: list[str] | None = None,
+) -> str:
+    turn = record.get("turn") if isinstance(record.get("turn"), dict) else {}
+    query = str(turn.get("query") or "")
+    context = record.get("context") if isinstance(record.get("context"), dict) else {}
+    prev_session = record.get("prev_session") if isinstance(record.get("prev_session"), list) else []
+    current_ts = optional_int_for_export(turn.get("timestamp"))
+    if current_ts is None:
+        source = record.get("source") if isinstance(record.get("source"), dict) else {}
+        current_ts = optional_int_for_export(source.get("timestamp"))
+    session, last_tts = training_history(
+        prev_session,
+        current_ts=current_ts,
+        session_num=session_num,
+        session_time_minutes=session_time_minutes,
+    )
+    fields = normalize_context_fields(context_fields)
+
+    instruction = "请参考用户的[当前query]、[对话历史]、[知识注入]、[系统状态]识别出[当前query]的[function]结果，[function]是python的code形式。\n"
+    instruction += "[知识注入]\n"
+    instruction += f"{context_prompt_block(context, fields)}\n"
+    instruction += "[系统状态]\n"
+    instruction += "{}\n"
+    instruction += "[对话历史]\n"
+    if session:
+        history_parts = [f"用户: {session_query}" for session_query in session]
+        if last_tts is not None:
+            history_parts.append(f"小爱: {last_tts}")
+        instruction += "\n".join(history_parts)
+        instruction += "\n"
+    instruction += "[当前query]\n"
+    instruction += f"用户: {query}\n"
+    instruction += "[function]\n"
+    return instruction
+
+
+def training_jsonl_line(
+    record: dict[str, Any],
+    *,
+    session_num: int,
+    session_time_minutes: int,
+    context_fields: list[str],
+    system_prompt: str,
+) -> str:
+    label = record.get("label") if isinstance(record.get("label"), dict) else {}
+    payload = {
+        "system": system_prompt,
+        "instruction": build_training_instruction(
+            record,
+            session_num=session_num,
+            session_time_minutes=session_time_minutes,
+            context_fields=context_fields,
+        ),
+        "output": str(label.get("target") or ""),
+    }
+    return json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
+
+
+def training_history(
+    prev_session: list[Any],
+    *,
+    current_ts: int | None,
+    session_num: int,
+    session_time_minutes: int,
+) -> tuple[list[str], str | None]:
+    if current_ts is None:
+        return [], None
+    session_time_ms = session_time_minutes * 60 * 1000
+    valid_items: list[tuple[int, str, str]] = []
+    for item in prev_session:
+        if not isinstance(item, dict):
+            continue
+        ts = optional_int_for_export(item.get("timestamp"))
+        query = str(item.get("query") or "")
+        tts = str(item.get("tts") or "")
+        if ts is not None and ts > 0 and query:
+            valid_items.append((ts, query, tts))
+    valid_items.sort(key=lambda item: item[0])
+
+    session: list[str] = []
+    last_tts: str | None = None
+    last_ts = current_ts
+    count = 0
+    for ts, query, tts in reversed(valid_items):
+        if last_ts - ts <= session_time_ms and last_ts >= ts:
+            session.append(query)
+            if count == 0:
+                last_tts = tts
+            last_ts = ts
+            count += 1
+            if count >= session_num:
+                break
+        else:
+            break
+    session.reverse()
+    return session, last_tts
+
+
+def context_prompt_block(context: dict[str, Any], fields: list[str]) -> str:
+    lines = ["{"]
+    for index, field in enumerate(fields):
+        comma = "," if index < len(fields) - 1 else ""
+        value = str(context.get(field, ""))
+        lines.append(f'"{field}": {json.dumps(value, ensure_ascii=False)}{comma}')
+    lines.append("}")
+    return "\n".join(lines)
+
+
+def normalize_context_fields(context_fields: list[str] | None) -> list[str]:
+    if context_fields is not None and not isinstance(context_fields, list):
+        raise ProductDataError("context_fields must be an array of strings")
+    fields = context_fields or list(DEFAULT_CONTEXT_FIELDS)
+    normalized = [field.strip() for field in fields if isinstance(field, str) and field.strip()]
+    return normalized or list(DEFAULT_CONTEXT_FIELDS)
+
+
+def optional_int_for_export(value: Any) -> int | None:
+    if isinstance(value, bool) or value is None or value == "":
+        return None
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        return None
+
+
+def category_label_from_target(target: str) -> str:
+    agent_match = re.match(r'''^Agent\s*\(\s*tag\s*=\s*["']([^"']+)["']\s*\)$''', target.strip())
+    if agent_match:
+        return agent_match.group(1)
+    function_names = re.findall(r"\b([A-Za-z_][A-Za-z0-9_]*)\s*\(", target)
+    if function_names:
+        return function_names[-1]
+    return target
+
+
 def dataset_table_columns() -> list[str]:
    return [
        "request_id",
@@ -524,6 +788,26 @@ def default_table_output_path(payload: dict[str, Any]) -> str:
    return str(Path("output") / "records.csv")


+def default_training_output_path(payload: dict[str, Any]) -> str:
+    if isinstance(payload.get("output_path"), str) and payload["output_path"].strip():
+        return payload["output_path"].strip()
+    if isinstance(payload.get("output_dir"), str) and payload["output_dir"].strip():
+        return str(Path(payload["output_dir"]).expanduser() / "training.jsonl")
+    if isinstance(payload.get("records_path"), str) and payload["records_path"].strip():
+        return str(Path(payload["records_path"]).expanduser().with_name("training.jsonl"))
+    return str(Path("output") / "training.jsonl")
+
+
+def default_planning_eval_output_path(payload: dict[str, Any]) -> str:
+    if isinstance(payload.get("output_path"), str) and payload["output_path"].strip():
+        return payload["output_path"].strip()
+    if isinstance(payload.get("output_dir"), str) and payload["output_dir"].strip():
+        return str(Path(payload["output_dir"]).expanduser() / "eval_planning.csv")
+    if isinstance(payload.get("records_path"), str) and payload["records_path"].strip():
+        return str(Path(payload["records_path"]).expanduser().with_name("eval_planning.csv"))
+    return str(Path("output") / "eval_planning.csv")
+
+
 def target_type(target: str) -> str:
    if re.match(r"^Agent\s*\(\s*tag\s*=", target):
        return "agent"
@@ -532,6 +816,13 @@ def target_type(target: str) -> str:
    return "unknown"


+def canonical_target(target: str) -> str:
+    agent_match = re.match(r'''^Agent\s*\(\s*tag\s*=\s*["']([^"']+)["']\s*\)$''', target)
+    if agent_match:
+        return f'Agent(tag="{agent_match.group(1)}")'
+    return target
+
+
 def record_id(*, source_type: str, batch_id: str, request_id: str, index: int) -> str:
    if source_type == "generated":
        return f"gen_{safe_token(batch_id)}_{index:06d}"