Add model labeling skill

2026-05-21 19:10:06 +08:00
parent 101a4a57c3
commit a3a4cdca14
4 changed files with 875 additions and 0 deletions
@@ -0,0 +1,600 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+"""批量调用线上模型 generate 接口给数据打标。
+
+输入通过 stdin 或 --input 传 JSON 对象，输出稳定 JSON 对象到 stdout。
+脚本只依赖 Python 标准库，便于在本地、Linux runtime 用户或远程工作区迁移执行。
+"""
+
+import argparse
+import csv
+import json
+import os
+import re
+import sys
+import time
+import urllib.error
+import urllib.request
+from io import StringIO
+from pathlib import Path
+from typing import Any
+
+
+DEFAULT_SYSTEM_PROMPT = "你是小爱同学，中文智能语音助手。"
+DEFAULT_PARAMETERS = {"max_new_tokens": 64}
+
+
+class LabelingError(ValueError):
+    """输入参数或数据格式错误。"""
+
+
+def main() -> int:
+    try:
+        payload = load_payload()
+        result = run(payload)
+        emit({"ok": True, **result})
+        return 0
+    except LabelingError as exc:
+        emit({"ok": False, "error": str(exc), **getattr(exc, "extra", {})})
+        return 1
+    except Exception as exc:  # noqa: BLE001 - CLI 需要稳定 JSON 错误
+        emit({"ok": False, "error": str(exc)})
+        return 1
+
+
+def load_payload() -> dict[str, Any]:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input", "-i", help="JSON 参数文件；不传则从 stdin 读取")
+    args = parser.parse_args()
+    text = Path(args.input).read_text(encoding="utf-8") if args.input else sys.stdin.read()
+    if not text.strip():
+        raise LabelingError("input JSON is required")
+    payload = json.loads(text)
+    if not isinstance(payload, dict):
+        raise LabelingError("input JSON must be an object")
+    return payload
+
+
+def emit(payload: dict[str, Any]) -> None:
+    print(json.dumps(payload, ensure_ascii=False, separators=(",", ":")))
+
+
+def run(payload: dict[str, Any]) -> dict[str, Any]:
+    input_path = require_string(payload, "input_path")
+    model_url = require_string(payload, "model_url")
+    dry_run = bool(payload.get("dry_run", False))
+    include_prompt = bool(payload.get("include_prompt", False))
+    max_records = optional_int(payload.get("max_records"))
+    start_index = int(payload.get("start_index") or 0)
+    timeout_seconds = float(payload.get("timeout_seconds") or 60)
+    output_path = str(payload.get("output_path") or "output/model_predictions.jsonl")
+    parameters = payload.get("parameters")
+    if parameters is None:
+        parameters = dict(DEFAULT_PARAMETERS)
+    if not isinstance(parameters, dict):
+        raise LabelingError("parameters must be an object")
+    headers = payload.get("headers")
+    if headers is None:
+        headers = {"Content-Type": "application/json"}
+    if not isinstance(headers, dict):
+        raise LabelingError("headers must be an object")
+    request_template = payload.get("request_template")
+    if request_template is not None and not isinstance(request_template, dict):
+        raise LabelingError("request_template must be an object")
+
+    source = read_input_file(input_path)
+    samples = normalize_samples(
+        source,
+        field_mapping=payload.get("field_mapping"),
+        system_prompt=str(payload.get("system_prompt") or DEFAULT_SYSTEM_PROMPT),
+        session_num=int(payload.get("session_num") or 5),
+        session_time_minutes=int(payload.get("session_time_minutes") or 5),
+    )
+    if start_index:
+        samples = samples[start_index:]
+    if max_records is not None:
+        samples = samples[:max_records]
+
+    summary = {
+        "input_path": input_path,
+        "source_format": source["format"],
+        "total_samples": len(samples),
+        "sample_preview": preview_samples(samples),
+    }
+    if dry_run:
+        return {**summary, "dry_run": True}
+
+    if not model_url.startswith(("http://", "https://")):
+        raise LabelingError("model_url must start with http:// or https://")
+    results: list[dict[str, Any]] = []
+    ok_count = 0
+    output = resolve_runtime_path(output_path)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    with output.open("w", encoding="utf-8") as file:
+        for sample in samples:
+            item = request_one(
+                sample,
+                model_url=model_url,
+                parameters=parameters,
+                headers={str(k): str(v) for k, v in headers.items()},
+                timeout_seconds=timeout_seconds,
+                request_template=request_template,
+                include_prompt=include_prompt,
+            )
+            if item.get("ok"):
+                ok_count += 1
+            results.append(item)
+            file.write(json.dumps(item, ensure_ascii=False, separators=(",", ":")) + "\n")
+
+    return {
+        **summary,
+        "dry_run": False,
+        "output_path": str(output),
+        "success_count": ok_count,
+        "failure_count": len(results) - ok_count,
+        "result_preview": results[:3],
+    }
+
+
+def request_one(
+    sample: dict[str, Any],
+    *,
+    model_url: str,
+    parameters: dict[str, Any],
+    headers: dict[str, str],
+    timeout_seconds: float,
+    request_template: dict[str, Any] | None,
+    include_prompt: bool,
+) -> dict[str, Any]:
+    prompt = str(sample["prompt"])
+    body = build_request_body(prompt, parameters, request_template)
+    started = time.time()
+    base = {
+        "index": sample["index"],
+        "request_id": sample.get("request_id", ""),
+        "query": sample.get("query", ""),
+        "gold_label": sample.get("gold_label", ""),
+        "source_format": sample.get("source_format", ""),
+    }
+    if include_prompt:
+        base["prompt"] = prompt
+    try:
+        request = urllib.request.Request(
+            model_url,
+            data=json.dumps(body, ensure_ascii=False).encode("utf-8"),
+            headers=headers,
+            method="POST",
+        )
+        with urllib.request.urlopen(request, timeout=timeout_seconds) as response:
+            text = response.read().decode("utf-8", errors="replace")
+            status = getattr(response, "status", 200)
+        latency_ms = int((time.time() - started) * 1000)
+        decoded = try_json(text)
+        return {
+            **base,
+            "prediction": extract_prediction(decoded, text),
+            "ok": 200 <= int(status) < 300,
+            "status": int(status),
+            "latency_ms": latency_ms,
+            "response": decoded if decoded is not None else text,
+        }
+    except urllib.error.HTTPError as exc:
+        text = exc.read().decode("utf-8", errors="replace")
+        return {
+            **base,
+            "prediction": extract_prediction(try_json(text), text),
+            "ok": False,
+            "status": exc.code,
+            "latency_ms": int((time.time() - started) * 1000),
+            "error": f"HTTP {exc.code}: {text[:500]}",
+        }
+    except Exception as exc:  # noqa: BLE001 - 单条失败不中断整体批次
+        return {
+            **base,
+            "prediction": "",
+            "ok": False,
+            "latency_ms": int((time.time() - started) * 1000),
+            "error": str(exc),
+        }
+
+
+def build_request_body(
+    prompt: str,
+    parameters: dict[str, Any],
+    request_template: dict[str, Any] | None,
+) -> dict[str, Any]:
+    if request_template is None:
+        return {"inputs": prompt, "parameters": parameters}
+    return replace_placeholders(request_template, {"prompt": prompt, "parameters": parameters})
+
+
+def replace_placeholders(value: Any, variables: dict[str, Any]) -> Any:
+    if isinstance(value, str):
+        if value == "{{prompt}}":
+            return variables["prompt"]
+        if value == "{{parameters}}":
+            return variables["parameters"]
+        return value.replace("{{prompt}}", str(variables["prompt"]))
+    if isinstance(value, list):
+        return [replace_placeholders(item, variables) for item in value]
+    if isinstance(value, dict):
+        return {key: replace_placeholders(item, variables) for key, item in value.items()}
+    return value
+
+
+def read_input_file(path: str) -> dict[str, Any]:
+    file_path = resolve_runtime_path(path)
+    if not file_path.exists():
+        raise LabelingError(f"input_path not found: {path}")
+    suffix = file_path.suffix.lower()
+    text = file_path.read_text(encoding="utf-8", errors="replace")
+    if suffix in {".csv", ".tsv"}:
+        delimiter = "\t" if suffix == ".tsv" else ","
+        rows = list(csv.DictReader(StringIO(text), delimiter=delimiter))
+        return {"format": "table", "path": str(file_path), "rows": rows, "columns": list(rows[0].keys()) if rows else []}
+    if suffix == ".jsonl":
+        rows = [json.loads(line) for line in text.splitlines() if line.strip()]
+        return {"format": "jsonl", "path": str(file_path), "rows": rows}
+    decoded = json.loads(text)
+    if isinstance(decoded, dict) and isinstance(decoded.get("records"), list):
+        decoded = decoded["records"]
+    if isinstance(decoded, list):
+        return {"format": "json", "path": str(file_path), "rows": decoded}
+    raise_with_mapping("JSON input must be an array or {records:[...]}", columns=list(decoded.keys()) if isinstance(decoded, dict) else [])
+
+
+def normalize_samples(
+    source: dict[str, Any],
+    *,
+    field_mapping: Any,
+    system_prompt: str,
+    session_num: int,
+    session_time_minutes: int,
+) -> list[dict[str, Any]]:
+    rows = source.get("rows")
+    if not isinstance(rows, list) or not rows:
+        raise LabelingError("input file contains no rows")
+    if isinstance(field_mapping, dict):
+        return samples_from_mapping(rows, field_mapping, source["format"], system_prompt)
+    first = rows[0]
+    if not isinstance(first, dict):
+        raise_with_mapping("rows must be objects")
+    if is_canonical_record(first):
+        return [
+            sample_from_record(index, row, system_prompt, session_num, session_time_minutes)
+            for index, row in enumerate(rows)
+            if isinstance(row, dict)
+        ]
+    if is_training_row(first):
+        return [sample_from_training(index, row, system_prompt) for index, row in enumerate(rows) if isinstance(row, dict)]
+    if source["format"] == "table":
+        columns = list(first.keys())
+        if "newPrompt" in columns:
+            return [sample_from_eval_row(index, row) for index, row in enumerate(rows) if isinstance(row, dict)]
+        if "query" in columns and "function" in columns:
+            return [sample_from_flow_row(index, row, system_prompt) for index, row in enumerate(rows) if isinstance(row, dict)]
+        raise_with_mapping("unrecognized table format", columns=columns)
+    raise_with_mapping("unrecognized JSON/JSONL format", columns=list(first.keys()))
+
+
+def sample_from_record(
+    index: int,
+    record: dict[str, Any],
+    system_prompt: str,
+    session_num: int,
+    session_time_minutes: int,
+) -> dict[str, Any]:
+    source = record.get("source") if isinstance(record.get("source"), dict) else {}
+    turn = record.get("turn") if isinstance(record.get("turn"), dict) else {}
+    return {
+        "index": index,
+        "request_id": str(source.get("request_id") or ""),
+        "query": str(turn.get("query") or ""),
+        "prompt": build_planning_prompt(record, system_prompt, session_num, session_time_minutes),
+        "gold_label": combined_label(record),
+        "source_format": "canonical_record_v1",
+    }
+
+
+def sample_from_training(index: int, row: dict[str, Any], system_prompt: str) -> dict[str, Any]:
+    prompt = str(row.get("prompt") or "")
+    if not prompt:
+        system = str(row.get("system") or system_prompt)
+        instruction = str(row.get("instruction") or row.get("input") or "")
+        prompt = wrap_chat_prompt(system, instruction)
+    return {
+        "index": index,
+        "request_id": str(row.get("request_id") or ""),
+        "query": extract_query_from_prompt(prompt),
+        "prompt": prompt,
+        "gold_label": str(row.get("output") or row.get("target") or ""),
+        "source_format": "training_jsonl",
+    }
+
+
+def sample_from_eval_row(index: int, row: dict[str, Any]) -> dict[str, Any]:
+    gold = str(row.get("code标签") or row.get("function") or row.get("target") or "")
+    complex_value = row.get("complex")
+    if complex_value not in (None, "") and not gold.startswith("complex="):
+        gold = f"complex={normalize_bool_literal(complex_value)}\n{gold}".rstrip()
+    prompt = str(row.get("newPrompt") or row.get("prompt") or "")
+    query = str(row.get("query") or "")
+    if not prompt:
+        prompt = build_planning_prompt(minimal_record(query, {}, [], gold), DEFAULT_SYSTEM_PROMPT, 5, 5)
+    return {
+        "index": index,
+        "request_id": str(row.get("request_id") or ""),
+        "query": query,
+        "prompt": prompt,
+        "gold_label": gold,
+        "source_format": "eval_csv",
+    }
+
+
+def sample_from_flow_row(index: int, row: dict[str, Any], system_prompt: str) -> dict[str, Any]:
+    record = record_from_flow_row(row)
+    return {
+        "index": index,
+        "request_id": str(row.get("request_id") or ""),
+        "query": str(row.get("query") or ""),
+        "prompt": build_planning_prompt(record, system_prompt, 5, 5),
+        "gold_label": str(row.get("function") or ""),
+        "source_format": "flow_csv",
+    }
+
+
+def samples_from_mapping(
+    rows: list[Any],
+    mapping: dict[str, Any],
+    source_format: str,
+    system_prompt: str,
+) -> list[dict[str, Any]]:
+    query_field = str(mapping.get("query") or "")
+    prompt_field = str(mapping.get("prompt") or "")
+    label_field = str(mapping.get("label") or mapping.get("target") or "")
+    request_id_field = str(mapping.get("request_id") or "")
+    if not query_field and not prompt_field:
+        raise LabelingError("field_mapping must provide query or prompt")
+    samples: list[dict[str, Any]] = []
+    for index, row in enumerate(rows):
+        if not isinstance(row, dict):
+            continue
+        prompt = str(row.get(prompt_field) or "")
+        query = str(row.get(query_field) or "")
+        if not prompt:
+            record = minimal_record(query, {}, [], str(row.get(label_field) or ""))
+            prompt = build_planning_prompt(record, system_prompt, 5, 5)
+        samples.append(
+            {
+                "index": index,
+                "request_id": str(row.get(request_id_field) or ""),
+                "query": query or extract_query_from_prompt(prompt),
+                "prompt": prompt,
+                "gold_label": str(row.get(label_field) or ""),
+                "source_format": f"{source_format}_mapped",
+            }
+        )
+    return samples
+
+
+def is_canonical_record(row: dict[str, Any]) -> bool:
+    return isinstance(row.get("turn"), dict) and bool(row["turn"].get("query"))
+
+
+def is_training_row(row: dict[str, Any]) -> bool:
+    return any(key in row for key in ("instruction", "system", "output", "prompt"))
+
+
+def build_planning_prompt(
+    record: dict[str, Any],
+    system_prompt: str,
+    session_num: int,
+    session_time_minutes: int,
+) -> str:
+    instruction = build_training_instruction(record, session_num, session_time_minutes)
+    return wrap_chat_prompt(system_prompt, instruction)
+
+
+def wrap_chat_prompt(system_prompt: str, instruction: str) -> str:
+    return (
+        f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
+        f"<|im_start|>user\n{instruction}<|im_end|>\n"
+        "<|im_start|>assistant\n"
+    )
+
+
+def build_training_instruction(record: dict[str, Any], session_num: int, session_time_minutes: int) -> 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(turn.get("timestamp"))
+    history = render_history(prev_session, current_ts, session_num, session_time_minutes)
+    return (
+        "请参考用户的[当前query]、[对话历史]、[知识注入]、[系统状态]识别出[当前query]的[function]结果，[function]是python的code形式。\n"
+        "[知识注入]\n"
+        f"{json.dumps({'location': str(context.get('location') or ''), 'rag': str(context.get('rag') or '')}, ensure_ascii=False, indent=0)}\n"
+        "[系统状态]\n"
+        "{}\n"
+        "[对话历史]\n"
+        f"{history}"
+        "[当前query]\n"
+        f"用户: {query}\n"
+        "[function]\n"
+    )
+
+
+def render_history(prev_session: list[Any], current_ts: int | None, session_num: int, session_time_minutes: int) -> str:
+    usable: list[dict[str, Any]] = []
+    for item in prev_session[-session_num:]:
+        if not isinstance(item, dict):
+            continue
+        ts = optional_int(item.get("timestamp"))
+        if current_ts is not None and ts is not None:
+            if abs(current_ts - ts) > session_time_minutes * 60_000:
+                continue
+        usable.append(item)
+    if not usable:
+        return ""
+    lines: list[str] = []
+    for item in usable:
+        query = str(item.get("query") or "").strip()
+        tts = str(item.get("tts") or "").strip()
+        if query:
+            lines.append(f"用户: {query}")
+        if tts:
+            lines.append(f"小爱: {tts}")
+    return "\n".join(lines) + ("\n" if lines else "")
+
+
+def record_from_flow_row(row: dict[str, Any]) -> dict[str, Any]:
+    prev_session = parse_json_cell(row.get("prev_session"), default=[])
+    context = parse_json_cell(row.get("context"), default={})
+    return minimal_record(
+        str(row.get("query") or ""),
+        context if isinstance(context, dict) else {},
+        prev_session if isinstance(prev_session, list) else [],
+        str(row.get("function") or ""),
+        request_id=str(row.get("request_id") or ""),
+        timestamp=optional_int(row.get("timestamp")),
+    )
+
+
+def minimal_record(
+    query: str,
+    context: dict[str, Any],
+    prev_session: list[Any],
+    target: str,
+    *,
+    request_id: str = "",
+    timestamp: int | None = None,
+) -> dict[str, Any]:
+    return {
+        "source": {"request_id": request_id, "timestamp": timestamp},
+        "turn": {"query": query, "timestamp": timestamp},
+        "prev_session": prev_session,
+        "context": context,
+        "label": {"target": target},
+        "dimensions": {},
+    }
+
+
+def combined_label(record: dict[str, Any]) -> str:
+    target = ""
+    label = record.get("label") if isinstance(record.get("label"), dict) else {}
+    if isinstance(label, dict):
+        target = str(label.get("target") or "")
+    dimensions = record.get("dimensions") if isinstance(record.get("dimensions"), dict) else {}
+    complex_value = dimensions.get("complex") if isinstance(dimensions, dict) else None
+    if isinstance(complex_value, bool):
+        return f"complex={'true' if complex_value else 'false'}\n{target}".rstrip()
+    return target
+
+
+def extract_query_from_prompt(prompt: str) -> str:
+    match = re.search(r"\[当前query\]\s*\n用户[:：]\s*(.+)", prompt)
+    return match.group(1).strip() if match else ""
+
+
+def normalize_bool_literal(value: Any) -> str:
+    text = str(value).strip().lower()
+    return "true" if text in {"true", "1", "yes", "y", "是", "复杂"} else "false"
+
+
+def parse_json_cell(value: Any, default: Any) -> Any:
+    if value in (None, ""):
+        return default
+    if isinstance(value, (dict, list)):
+        return value
+    try:
+        return json.loads(str(value))
+    except json.JSONDecodeError:
+        return default
+
+
+def resolve_runtime_path(path: str) -> Path:
+    raw = Path(path).expanduser()
+    if raw.is_absolute():
+        return raw
+    scratchpad = Path(str(Path.cwd()))
+    if os.environ.get("PYTHON_EXEC_SCRATCHPAD"):
+        scratchpad = Path(os.environ["PYTHON_EXEC_SCRATCHPAD"]).expanduser()
+    parts = raw.parts
+    if parts and parts[0] in {"output", "outputs"}:
+        return scratchpad.parent / "output" / Path(*parts[1:])
+    if parts and parts[0] in {"input", "inputs"}:
+        return scratchpad.parent / "input" / Path(*parts[1:])
+    if parts and parts[0] in {"scratchpad", "scratch"}:
+        return scratchpad / Path(*parts[1:])
+    return raw
+
+
+def optional_int(value: Any) -> int | None:
+    if value in (None, ""):
+        return None
+    try:
+        return int(float(str(value)))
+    except (TypeError, ValueError):
+        return None
+
+
+def extract_prediction(decoded: Any, text: str) -> str:
+    if isinstance(decoded, dict):
+        for key in ("generated_text", "text", "output", "response", "result"):
+            value = decoded.get(key)
+            if isinstance(value, str):
+                return value.strip()
+        outputs = decoded.get("outputs")
+        if isinstance(outputs, list) and outputs:
+            first = outputs[0]
+            if isinstance(first, str):
+                return first.strip()
+            if isinstance(first, dict):
+                return extract_prediction(first, json.dumps(first, ensure_ascii=False))
+        choices = decoded.get("choices")
+        if isinstance(choices, list) and choices:
+            first = choices[0]
+            if isinstance(first, dict):
+                message = first.get("message")
+                if isinstance(message, dict) and isinstance(message.get("content"), str):
+                    return message["content"].strip()
+                if isinstance(first.get("text"), str):
+                    return first["text"].strip()
+    return text.strip()
+
+
+def try_json(text: str) -> Any:
+    try:
+        return json.loads(text)
+    except json.JSONDecodeError:
+        return None
+
+
+def preview_samples(samples: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    return [
+        {
+            "index": item.get("index"),
+            "query": item.get("query"),
+            "gold_label": item.get("gold_label"),
+            "prompt_preview": str(item.get("prompt") or "")[:200],
+        }
+        for item in samples[:3]
+    ]
+
+
+def require_string(payload: dict[str, Any], key: str) -> str:
+    value = payload.get(key)
+    if not isinstance(value, str) or not value.strip():
+        raise LabelingError(f"{key} is required")
+    return value.strip()
+
+
+def raise_with_mapping(message: str, columns: list[str] | None = None) -> None:
+    exc = LabelingError(message)
+    exc.extra = {"needs_mapping": True, "columns": columns or []}  # type: ignore[attr-defined]
+    raise exc
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())