Files
Video-Tree-TRM5/app/harness/inference.py
T

591 lines
21 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""async 推理编排 — 训练循环的 forward()。
从 TRM4 core/harness/inference.py (~560 行) 迁移,重大重构:
- 同步 ThreadPoolExecutor → asyncio.Semaphore + asyncio.gather
- LLMClient.from_env() 每题构造 → llm: LLMProvider 注入共享
- SentenceTransformer/OCR 内部构造 → 调用方通过 tool_dispatch_fn 注入
- run_id 必传,空串 → ValueError
- _aggregate_results 从内存 results 聚合(非 DB 回读)
- record_run 由调用方(Runner)负责
- prompt 构建由调用方注入 prompt_builder
"""
from __future__ import annotations
import asyncio
import json
import sqlite3
from collections import defaultdict
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
from loguru import logger
from app.harness.question_units import build_units, unit_correctness
from core.agent.loop import AgentLoop
if TYPE_CHECKING:
from collections.abc import Callable
from app.harness.log import HarnessLog
from core.agent.types import LoopResult
from core.protocols import LLMProvider
from core.types import GeneratedQuestion, QuestionUnit
@dataclass(frozen=True)
class InferenceResult:
"""推理聚合结果(正确率按 unit 粒度)。
属性:
run_id: 运行标识。
accuracy: unit 级正确率(correct / total)。
total: unit 总数(single 数 + pair 数,孤儿 pair 已剔除不计入)。
correct: 正确 unit 数(single 单题正确;pair 走 original/mirror 双向 AND)。
per_task_type: 按题型分组的 unit 级指标 {task_type: {accuracy, total, correct}}。
steps_mean: 平均步数(record 粒度,逐题溯源)。
token_usage: token 总用量 {prompt_tokens, completion_tokens}record 粒度)。
stop_reason_counts: 终止原因计数 {reason: count}record 粒度)。
"""
run_id: str
accuracy: float
total: int
correct: int
per_task_type: dict[str, dict]
steps_mean: float
token_usage: dict[str, int]
stop_reason_counts: dict[str, int]
# ---------------------------------------------------------------------------
# 表 Schema 定义(5 张表,保留 TRM4 全部 schema
# ---------------------------------------------------------------------------
PREDICTIONS_SCHEMA: dict[str, str] = {
"video_id": "TEXT",
"question_id": "TEXT",
"task_type": "TEXT",
"prediction": "TEXT",
"answer": "TEXT",
"evidence": "TEXT",
"reasoning": "TEXT",
"steps_used": "INTEGER",
"prompt_tokens": "INTEGER",
"completion_tokens": "INTEGER",
"stop_reason": "TEXT",
"steps_json": "JSON",
}
TRACES_SCHEMA: dict[str, str] = {
"video_id": "TEXT",
"question_id": "TEXT",
"step": "INTEGER",
"tool_name": "TEXT",
"tool_args": "JSON",
"tool_output": "TEXT",
"thought": "TEXT",
}
VALIDATION_FLAGS_SCHEMA: dict[str, str] = {
"video_id": "TEXT",
"question_id": "TEXT",
"has_l3_visit": "INTEGER",
"l1_count": "INTEGER",
"l2_count": "INTEGER",
"l3_count": "INTEGER",
}
ANCHOR_CHECK_SCHEMA: dict[str, str] = {
"video_id": "TEXT",
"question_id": "TEXT",
"step": "INTEGER",
"n_assertions": "INTEGER",
"n_anchored": "INTEGER",
"n_illegal": "INTEGER",
"n_expanded": "INTEGER",
"n_trunc": "INTEGER",
"output_chars": "INTEGER",
}
OF_HEALTH_SCHEMA: dict[str, str] = {
"video_id": "TEXT",
"question_id": "TEXT",
"step": "INTEGER",
"ocr_injected": "INTEGER",
"ocr_chars": "INTEGER",
"ocr_failed": "INTEGER",
"discrepancy": "INTEGER",
"abstain": "INTEGER",
}
# ---------------------------------------------------------------------------
# 内部工具
# ---------------------------------------------------------------------------
class _DispatcherAdapter:
"""将裸 async callable 包装为 ToolDispatcher Protocol 实例。
AgentLoop 要求 ToolDispatcher(有 dispatch 方法),而 run_inference
接收的 tool_dispatch_fn 是裸 async callable。此适配器桥接两者。
参数:
fn: async def (tool_name, args, *, context) -> str。
"""
def __init__(self, fn: Callable[..., Any]) -> None:
self._fn = fn
async def dispatch(
self, tool_name: str, args: dict[str, Any], *, context: dict[str, Any]
) -> str:
"""转发工具调用给被包装的 callable。"""
return await self._fn(tool_name, args, context=context)
def _to_text_field(value: Any) -> str:
"""把 prediction 的 evidence/reasoning 归一为可入库的文本。
LLM 有时把这些字段返回成 list 或 dict(而非字符串)。sqlite 无法绑定
非标量类型,直接入库会抛 ProgrammingError 致该题丢失预测行、进而触发
rollout 完整性护栏中止整轮。凡非 str 一律 JSON 序列化为文本。
参数:
value: evidence/reasoning 原始值(可能是 str/list/dict)。
返回:
可直接入库的字符串。
"""
if isinstance(value, str):
return value
return json.dumps(value, ensure_ascii=False)
def _normalize_prediction(answer: object) -> str | None:
"""归一化 prediction 落库值。
LLM 提交的 answer 有时是 list/dict(如 {'answer': ['B']}),sqlite 无法绑定
非标量类型直接入库会抛 ProgrammingError 击穿整轮 gather。None 保留(INFRA 空
预测语义,供正确率判定天然计错);str 原样;其余 JSON 序列化为文本。
参数:
answer: LoopResult.result 中的 answer 原始值(可能是 None/str/list/dict)。
返回:
None(保留空预测语义)或可直接入库的字符串。
"""
if answer is None or isinstance(answer, str):
return answer
return _to_text_field(answer)
def _zero_result(run_id: str) -> InferenceResult:
"""空记录时的零值 InferenceResult。
参数:
run_id: 运行标识。
返回:
全零的 InferenceResult。
"""
return InferenceResult(
run_id=run_id,
accuracy=0.0,
total=0,
correct=0,
per_task_type={},
steps_mean=0.0,
token_usage={"prompt_tokens": 0, "completion_tokens": 0},
stop_reason_counts={},
)
def _group_by_task_type(graded: list[tuple[QuestionUnit, bool]]) -> dict[str, dict[str, Any]]:
"""按 task_type 分组聚合 unit 级正确率指标。
pair 单元整体计 1 个 unit,归入其 task_typesingle 单元计 1 个 unit。
参数:
graded: (单元, 该单元是否整体正确) 元组列表。
返回:
{task_type: {accuracy, total, correct}} 映射(unit 粒度)。
"""
task_groups: dict[str, list[bool]] = defaultdict(list)
for unit, is_correct in graded:
task_groups[unit.task_type].append(is_correct)
per_task_type: dict[str, dict[str, Any]] = {}
for task_type, verdicts in task_groups.items():
t_total = len(verdicts)
t_correct = sum(verdicts)
per_task_type[task_type] = {
"accuracy": t_correct / t_total,
"total": t_total,
"correct": t_correct,
}
return per_task_type
def _is_valid_pair(group: list[GeneratedQuestion]) -> bool:
"""判定同一 pair_id 分组是否为合法孪生对(恰好 1 original + 1 mirror,无多余)。
要求分组总数恰为 2 且角色齐备唯一;有额外非法 role 记录(total>2)或角色
缺失/重复均视为非法,交由调用方剔除,防非法记录混入 build_units。
参数:
group: 归属同一 pair_id 的题目列表。
返回:
合法孪生对为 True,否则 False。
"""
if len(group) != 2:
return False
originals = sum(1 for q in group if q.question_role == "pair_original")
mirrors = sum(1 for q in group if q.question_role == "pair_mirror")
return originals == 1 and mirrors == 1
def _drop_orphan_pairs(questions: list[GeneratedQuestion]) -> list[GeneratedQuestion]:
"""剔除收不齐 2 条 / 角色非法的孤儿 pair,告警不静默。
每条题目均会各答一次并逐题落库;能否合成 pair 单元仅取决于 questions
是否同时含该 pair_id 的 original + mirror(且无多余非法记录)。非法者告警并
整对剔除,使后续 build_units 只面对合法孪生对(不触发 fail-fast),孤儿 unit
不计入 total(对齐设计 §8 聚合入口的"告警 + 剔除")。
参数:
questions: 待聚合的题目列表(可混含 single 与孪生对成员)。
返回:
可安全交给 build_units 的题目列表(single 全保留,pair 仅保留合法成对者)。
"""
by_pair: dict[str, list[GeneratedQuestion]] = defaultdict(list)
singles: list[GeneratedQuestion] = []
for q in questions:
if q.pair_id:
by_pair[q.pair_id].append(q)
else:
singles.append(q)
kept_pairs: list[GeneratedQuestion] = []
for pair_id, group in by_pair.items():
if _is_valid_pair(group):
kept_pairs.extend(group)
else:
logger.warning(
"孤儿 pair {}:非法配对(total={}),剔除该 unit 不计入 total",
pair_id,
len(group),
)
return singles + kept_pairs
def _per_question_correctness(records: list[dict[str, Any]]) -> dict[str, bool]:
"""由逐题 record 构造 question_id → 该题作答是否正确 的映射。
prediction 为 None(作答异常)时与 answer 不相等 → False,天然计错。
参数:
records: _run_single_question 返回的逐题 record 列表。
返回:
{question_id: prediction == answer} 映射,供 unit_correctness 取值。
"""
return {r["question_id"]: r["prediction"] == r["answer"] for r in records}
def _grade_unit(unit: QuestionUnit, per_q: dict[str, bool]) -> bool:
"""判定单元整体正确性,缺 prediction 时 fail-loud(带上下文)。
_drop_orphan_pairs 已剔除孤儿/非法配对,正常情况下 unit 内每题都应有对应
record;若仍缺失说明聚合不变量被破坏(如 records 与 questions 不同源)。此处
显式抛带上下文的 ValueErrorfail-loud,不 catch/不跳过/不兜底),而非放任
unit_correctness 抛裸 KeyError 丢失定位信息。
参数:
unit: 待判定单元。
per_q: question_id → 该题是否作答正确 的映射。
返回:
单元整体是否正确(single 即单题正确;pair 走双向 AND)。
异常:
ValueError: unit 内某 question_id 不在 per_q 中(聚合不变量被破坏)。
"""
missing = [q.question_id for q in unit.questions if q.question_id not in per_q]
if missing:
raise ValueError(
f"unit {unit.unit_id} 的 question {missing} 缺 prediction"
"_drop_orphan_pairs 后不应发生,聚合不变量被破坏)"
)
return unit_correctness(unit, per_q)
def _aggregate_results(
records: list[dict[str, Any]],
questions: list[GeneratedQuestion],
run_id: str,
) -> InferenceResult:
"""从内存 records + 题目列表按 unit 粒度聚合推理指标。
逐题 record 保留逐题溯源(token/steps/stop_reason 诊断仍按 record 汇总);
正确率则按 unit 粒度计:single 计 1AR pair 经 build_units 收齐 original +
mirror 后走 unit_correctness 的双向 AND 判定,整对计 1 个 unit。孤儿 pair
在 _drop_orphan_pairs 中告警 + 剔除,不计入 total。
参数:
records: _run_single_question 返回的逐题 record 列表。
questions: 与 records 对应的题目列表(提供 pair_id/question_role 元数据)。
run_id: 当前运行标识。
返回:
InferenceResult 冻结实例(total/correct/per_task_type 为 unit 粒度)。
"""
if not records:
return _zero_result(run_id)
per_q = _per_question_correctness(records)
units = build_units(_drop_orphan_pairs(questions))
graded = [(unit, _grade_unit(unit, per_q)) for unit in units]
total = len(graded)
correct = sum(1 for _, is_correct in graded if is_correct)
stop_counts: dict[str, int] = defaultdict(int)
for r in records:
stop_counts[r["stop_reason"]] += 1
n_records = len(records)
return InferenceResult(
run_id=run_id,
accuracy=correct / total if total else 0.0,
total=total,
correct=correct,
per_task_type=_group_by_task_type(graded),
steps_mean=sum(r["steps_used"] for r in records) / n_records,
token_usage={
"prompt_tokens": sum(r["prompt_tokens"] for r in records),
"completion_tokens": sum(r["completion_tokens"] for r in records),
},
stop_reason_counts=dict(stop_counts),
)
# ---------------------------------------------------------------------------
# 单题推理
# ---------------------------------------------------------------------------
async def _run_single_question(
qa: GeneratedQuestion,
*,
llm: LLMProvider,
tool_dispatch_fn: Callable[..., Any],
prompt_builder: Callable[[GeneratedQuestion], tuple[str, str]],
log: HarnessLog,
max_steps: int,
plugins: list[object],
run_id: str,
) -> dict[str, Any]:
"""执行单道题目的 Agent 推理。
悲观默认值:record 初始 stop_reason="error",成功后覆盖。
prediction 必落库:log.insert 在 try/except 之后(无论成败)。
参数:
qa: 待推理的题目。
llm: LLMProvider 共享实例。
tool_dispatch_fn: async 工具调度函数 (tool_name, args, *, context) -> str。
prompt_builder: (GeneratedQuestion) -> (system_prompt, user_prompt)。
log: HarnessLog 实例(线程安全)。
max_steps: AgentLoop 最大步数。
plugins: pluggy 插件列表。
run_id: 运行标识,用作 cache_salt——run_id 含 _e{epoch} 天然跨 epoch 重采样、
同 epoch 续跑命中缓存(算法 #10 透传)。
返回:
预测结果字典(含 video_id, question_id, prediction, answer 等)。
"""
# run_id 必须显式入 record:HarnessLog.insert 缺省用**实例** run_id 填充,
# 连续并发 gate 共享单一 gate_log(实例 run_id 为 step 级)时,各臂行必须
# 落自己的臂 run_id,否则 validate 回读 _load_run_rows(臂 run_id) 为空。
record: dict[str, Any] = {
"run_id": run_id,
"video_id": qa.video_id,
"question_id": qa.question_id,
"task_type": qa.task_type,
"prediction": None,
"answer": qa.answer,
"evidence": "",
"reasoning": "",
"steps_used": 0,
"prompt_tokens": 0,
"completion_tokens": 0,
"stop_reason": "error", # 悲观默认
"steps_json": "[]",
}
try:
system_prompt, user_prompt = prompt_builder(qa)
dispatcher = _DispatcherAdapter(tool_dispatch_fn)
loop = AgentLoop(llm, max_steps=max_steps)
loop_result: LoopResult = await loop.run(
system_prompt,
user_prompt,
dispatcher,
plugins=plugins,
session_id=qa.question_id,
cache_salt=run_id,
)
result_dict = loop_result.result if isinstance(loop_result.result, dict) else {}
evidence = _to_text_field(result_dict.get("evidence", ""))
reasoning = _to_text_field(result_dict.get("reasoning", ""))
record.update(
{
"prediction": _normalize_prediction(result_dict.get("answer")),
"evidence": evidence,
"reasoning": reasoning,
"steps_used": loop_result.steps_used,
"prompt_tokens": loop_result.token_usage["prompt_tokens"],
"completion_tokens": loop_result.token_usage["completion_tokens"],
"stop_reason": loop_result.stop_reason,
"steps_json": json.dumps(
[
{
"thought": s.thought,
"tool_call": s.tool_call,
"tool_output": s.tool_output,
}
for s in loop_result.steps
],
ensure_ascii=False,
),
}
)
except Exception:
logger.exception("[{}] QA {} 执行异常", qa.video_id, qa.question_id)
# prediction 必落库(try 外,无论成败);绑定异常降级为最小 error 行,不击穿 gather
try:
await asyncio.to_thread(log.insert, "predictions", record)
except (sqlite3.InterfaceError, sqlite3.ProgrammingError):
logger.exception("[{}] QA {} 落库绑定异常,降级为 error 行", qa.video_id, qa.question_id)
record["prediction"] = None
record["stop_reason"] = "error"
await asyncio.to_thread(
log.insert,
"predictions",
{k: v for k, v in record.items() if isinstance(v, (str, int, float, type(None)))},
)
return record
# ---------------------------------------------------------------------------
# 建表
# ---------------------------------------------------------------------------
def _ensure_tables(log: HarnessLog) -> None:
"""创建推理所需的 5 张表。
参数:
log: HarnessLog 实例。
"""
log.create_table("predictions", PREDICTIONS_SCHEMA)
log.create_table("traces", TRACES_SCHEMA)
log.create_table("validation_flags", VALIDATION_FLAGS_SCHEMA)
log.create_table("anchor_check", ANCHOR_CHECK_SCHEMA)
log.create_table("observe_frame_health", OF_HEALTH_SCHEMA)
# ---------------------------------------------------------------------------
# 公共入口
# ---------------------------------------------------------------------------
async def run_inference(
questions: list[GeneratedQuestion],
*,
llm: LLMProvider,
tool_dispatch_fn: Callable[..., Any],
prompt_builder: Callable[[GeneratedQuestion], tuple[str, str]],
log: HarnessLog,
run_id: str,
concurrency: int,
max_steps: int,
skill_mode: str,
plugins_factory: Callable[[str, str], list[object]] | None = None,
) -> InferenceResult:
"""在视频树上执行 Agent 推理,对应训练循环的 forward()。
参数:
questions: 待推理的题目列表。
llm: LLMProvider 共享实例(依赖注入)。
tool_dispatch_fn: async 工具调度函数 (tool_name, args, *, context) -> str。
prompt_builder: prompt 构建函数 (GeneratedQuestion) -> (system_prompt, user_prompt)。
log: HarnessLog 实例(由调用方管理生命周期)。
run_id: 运行标识(必传,空串 → ValueError)。
concurrency: 最大并发数(asyncio.Semaphore 控制)。
max_steps: AgentLoop 单题最大步数。
skill_mode: "auto" / "manual" / "none"(传递给调用方的 prompt/plugin 构建逻辑)。
plugins_factory: 可选的插件工厂 (video_id, question_id) -> plugins 列表。
返回:
InferenceResult(含 accuracy、per_task_type 等聚合指标)。
异常:
ValueError: run_id 为空串或纯空白。
"""
if not run_id or not run_id.strip():
raise ValueError("run_id 不得为空串或纯空白")
_ensure_tables(log)
if not questions:
logger.info("题目列表为空,返回零值 InferenceResult")
return _aggregate_results([], [], run_id)
sem = asyncio.Semaphore(concurrency)
total_count = len(questions)
async def _bounded(index: int, qa: GeneratedQuestion) -> dict[str, Any]:
"""信号量限流的单题推理包装。"""
async with sem:
plugins = (
plugins_factory(qa.video_id, qa.question_id) if plugins_factory is not None else []
)
result = await _run_single_question(
qa,
llm=llm,
tool_dispatch_fn=tool_dispatch_fn,
prompt_builder=prompt_builder,
log=log,
max_steps=max_steps,
plugins=plugins,
run_id=run_id,
)
logger.info(
"[{}/{}] {} QA {} 完成 (stop={})",
index + 1,
total_count,
qa.video_id,
qa.question_id,
result["stop_reason"],
)
return result
results = await asyncio.gather(*[_bounded(i, qa) for i, qa in enumerate(questions)])
inference_result = _aggregate_results(list(results), questions, run_id)
logger.info(
"推理完成: accuracy={:.2%} ({}/{})",
inference_result.accuracy,
inference_result.correct,
inference_result.total,
)
return inference_result