初始化项目,由ModelHub XC社区提供模型

Model: tawer12/qwen3-4b-recursive-sft-v3.2
Source: Original Platform
This commit is contained in:
ModelHub XC
2026-10-22 13:37:40 +08:00
commit bf506ebcbe
18 changed files with 153271 additions and 0 deletions

570
recursive_engine.py Normal file
View File

@@ -0,0 +1,570 @@
#!/usr/bin/env python3
"""Recursive inference engine for stateless recursive math-solving models."""
from __future__ import annotations
import json
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
TAG_RE = re.compile(r"</?(problem|action|parallel|path)>", re.IGNORECASE)
NUMBER_RE = re.compile(r"-?\d+(?:,\d{3})*(?:\.\d+)?|-?\d+/\d+")
FINAL_ANSWER_LINE_RE = re.compile(
r"(?im)^[ \t]*Final[ \t]+answer[ \t]*:[ \t]*(?P<answer>[^\n]+)$"
)
DEFAULT_PREFIX_PROMPT_FILE = Path("prompts") / "recursive_sft_prefix_prompt.txt"
FORCED_ROOT_OUTPUT_PREFIXES = {
"SOLVE_DIRECTLY": "<action>\nSOLVE_DIRECTLY\n",
"DECOMPOSE": "<action>\nDECOMPOSE\n</action>\n<parallel>\n",
}
@dataclass
class ModelCall:
call_id: int
depth: int
role: str
problem: str
input_text: str
output_text: str
input_tokens: int
output_tokens: int
action_type: str | None = None
parse_error: str | None = None
forced_output_prefix: str = ""
@dataclass
class RolloutResult:
problem: str
final_answer: str | None
full_trace: str
calls: list[ModelCall] = field(default_factory=list)
parse_errors: list[str] = field(default_factory=list)
critical_path_tokens: int = 0
stopped_reason: str = "completed"
@property
def parsed_ok(self) -> bool:
return not self.parse_errors
@property
def model_call_count(self) -> int:
return len(self.calls)
@dataclass
class NodeResult:
final_answer: str | None
trace: str
critical_path_tokens: int
parse_errors: list[str]
stopped_reason: str = "completed"
def load_prefix_template(path: str | Path | None = None) -> str:
prompt_path = Path(path) if path is not None else DEFAULT_PREFIX_PROMPT_FILE
return prompt_path.read_text(encoding="utf-8")
def format_model_input(problem: str, prefix_template: str) -> str:
return prefix_template.strip().replace("{PROBLEM}", problem.strip())
def token_count(tokenizer: Any, text: str) -> int:
return len(tokenizer.encode(text, add_special_tokens=False))
def extract_blocks(text: str, tag: str) -> list[str]:
return [
match.group(1).strip()
for match in re.finditer(rf"<{tag}>(.*?)</{tag}>", text, flags=re.DOTALL | re.IGNORECASE)
]
def normalize_answer(value: Any) -> str | None:
if value is None:
return None
text = str(value).strip()
if not text:
return None
if "####" in text:
text = text.split("####")[-1].strip()
boxed = re.findall(r"\\boxed\{([^{}]+)\}", text)
if boxed:
text = boxed[-1]
final = re.findall(r"(?i)Final\s+answer\s*:\s*([^\n<]+)", text)
if final:
text = final[-1]
text = text.replace("$", "").replace(",", "")
text = re.sub(r"\\(?:text|mathrm)\{([^{}]+)\}", r"\1", text)
text = re.sub(r"\s+", "", text)
text = text.rstrip(".")
return text or None
def extract_final_answer(text: str) -> str | None:
patterns = [
r"(?i)Final\s+answer\s*:\s*([^\n<]+)",
r"(?i)answer\s+is\s*([^\n<]+)",
r"\\boxed\{([^{}]+)\}",
]
for pattern in patterns:
matches = re.findall(pattern, text)
if matches:
return normalize_answer(matches[-1])
matches = NUMBER_RE.findall(text)
if matches:
return normalize_answer(matches[-1])
return None
def answers_match(pred: str | None, gt: str | None) -> bool:
pred_norm = normalize_answer(pred)
gt_norm = normalize_answer(gt)
if pred_norm is None or gt_norm is None:
return False
if pred_norm == gt_norm:
return True
try:
from math_verify import LatexExtractionConfig, parse, verify
pred_parsed = parse(pred_norm, extraction_config=[LatexExtractionConfig()])
gt_parsed = parse(gt_norm, extraction_config=[LatexExtractionConfig()])
if pred_parsed and gt_parsed and verify(gt_parsed, pred_parsed):
return True
except Exception:
pass
return pred_norm.lower() == gt_norm.lower()
def find_matching_tag(text: str, tag: str, open_start: int, open_end: int) -> tuple[int, int] | None:
depth = 1
for match in TAG_RE.finditer(text, open_end):
token = match.group(0).lower()
name = match.group(1).lower()
if name != tag:
continue
if token.startswith("</"):
depth -= 1
if depth == 0:
return match.start(), match.end()
else:
depth += 1
return None
def parse_model_output(output: str) -> tuple[str, list[str]]:
action_match = re.search(r"<action>(.*?)</action>", output, flags=re.DOTALL | re.IGNORECASE)
if not action_match:
raise ValueError("missing <action>...</action> block")
action_body = action_match.group(1).strip()
first_line = action_body.splitlines()[0].strip().upper() if action_body else ""
if first_line.startswith("SOLVE_DIRECTLY"):
return "SOLVE_DIRECTLY", []
if first_line.startswith("AGGREGATE"):
return "AGGREGATE", []
if not first_line.startswith("DECOMPOSE"):
raise ValueError(f"unknown action type in first action line: {first_line!r}")
next_nonspace = re.search(r"\S", output[action_match.end() :])
if not next_nonspace:
raise ValueError("DECOMPOSE action is not followed by <parallel>")
parallel_start = action_match.end() + next_nonspace.start()
if not output[parallel_start:].lower().startswith("<parallel>"):
raise ValueError("DECOMPOSE action is not immediately followed by <parallel>")
parallel_open_end = parallel_start + len("<parallel>")
close = find_matching_tag(output, "parallel", parallel_start, parallel_open_end)
if close is None:
raise ValueError("unclosed <parallel> block")
parallel_body = output[parallel_open_end : close[0]]
paths = extract_blocks(f"<parallel>{parallel_body}</parallel>", "path")
if len(paths) < 2:
raise ValueError(f"DECOMPOSE parallel block has {len(paths)} path(s), expected at least 2")
if any(not path for path in paths):
raise ValueError("DECOMPOSE parallel block contains an empty path")
return "DECOMPOSE", paths
def validate_action_protocol(output: str, action_type: str, role: str) -> None:
"""Enforce role-specific actions and explicit answers for terminal calls."""
normalized_role = role.strip().lower()
if normalized_role == "root":
allowed_actions = {"DECOMPOSE", "SOLVE_DIRECTLY"}
elif normalized_role == "aggregate":
allowed_actions = {"AGGREGATE"}
elif normalized_role == "child" or normalized_role.startswith("path_"):
allowed_actions = {"DECOMPOSE", "SOLVE_DIRECTLY"}
else:
raise ValueError(f"unknown call role: {role!r}")
if action_type not in allowed_actions:
expected = " or ".join(sorted(allowed_actions))
raise ValueError(
f"{normalized_role} call returned {action_type}, expected {expected}"
)
if action_type not in {"SOLVE_DIRECTLY", "AGGREGATE"}:
return
action_match = re.search(
r"<action>(.*?)</action>", output, flags=re.DOTALL | re.IGNORECASE
)
action_body = action_match.group(1) if action_match else ""
final_answer = FINAL_ANSWER_LINE_RE.search(action_body)
if final_answer is None or not final_answer.group("answer").strip():
raise ValueError(f"{action_type} action is missing a non-empty Final answer: line")
def extract_final_answer_display(text: str) -> str | None:
patterns = [
r"(?i)Final\s+answer\s*:\s*([^\n<]+)",
r"\\boxed\{([^{}]+)\}",
]
for pattern in patterns:
matches = re.findall(pattern, text)
if matches:
return matches[-1].strip().rstrip(".")
return None
def child_result_summary(result: NodeResult) -> str:
display = extract_final_answer_display(result.trace)
if display is not None:
return display
if result.final_answer is not None:
return str(result.final_answer)
fallback = extract_final_answer(result.trace)
if fallback is not None:
return str(fallback)
return "PARSE_ERROR" if result.parse_errors else "UNKNOWN"
def build_aggregation_problem(parent_problem: str, paths: list[str], child_results: list[NodeResult]) -> str:
lines = [
f"The original problem is: {parent_problem.strip()}",
"The paths spawned were:",
]
for idx, (path, result) in enumerate(zip(paths, child_results), start=1):
summary = child_result_summary(result)
lines.append(f"{idx}. {path.strip()} (Result: {summary}).")
lines.append("Aggregate these results to find the answer to the original problem.")
return "\n".join(lines)
class RecursiveEngine:
def __init__(
self,
*,
model: Any,
tokenizer: Any,
prefix_template: str,
max_new_tokens: int = 768,
temperature: float = 0.0,
top_p: float = 1.0,
top_k: int | None = None,
max_depth: int = 3,
max_calls: int = 10,
force_root_action: str | None = None,
) -> None:
self.model = model
self.tokenizer = tokenizer
self.prefix_template = prefix_template
self.max_new_tokens = max_new_tokens
self.temperature = temperature
self.top_p = top_p
self.top_k = top_k
self.max_depth = max_depth
self.max_calls = max_calls
normalized_force = force_root_action.upper() if force_root_action else None
if normalized_force not in {None, *FORCED_ROOT_OUTPUT_PREFIXES}:
raise ValueError(
"force_root_action must be SOLVE_DIRECTLY, DECOMPOSE, or None"
)
self.force_root_action = normalized_force
self.calls: list[ModelCall] = []
def generate_once(self, input_text: str, *, output_prefix: str = "") -> str:
import torch
generation_input = input_text + output_prefix
inputs = self.tokenizer(generation_input, return_tensors="pt")
device = next(self.model.parameters()).device
inputs = {key: value.to(device) for key, value in inputs.items()}
generation_kwargs = {
"max_new_tokens": self.max_new_tokens,
"do_sample": self.temperature > 0,
"top_p": self.top_p,
"use_cache": True,
"pad_token_id": self.tokenizer.pad_token_id or self.tokenizer.eos_token_id,
"eos_token_id": self.tokenizer.eos_token_id,
}
if self.temperature > 0:
generation_kwargs["temperature"] = self.temperature
if self.top_k is not None:
generation_kwargs["top_k"] = self.top_k
with torch.no_grad():
generated = self.model.generate(
**inputs,
**generation_kwargs,
)
new_tokens = generated[0, inputs["input_ids"].shape[1] :]
continuation = self.tokenizer.decode(new_tokens, skip_special_tokens=False).strip()
return output_prefix + continuation
def rollout(self, problem: str) -> RolloutResult:
self.calls = []
root = self._solve_node(problem=problem, depth=0, role="root")
return RolloutResult(
problem=problem,
final_answer=root.final_answer,
full_trace=root.trace,
calls=self.calls,
parse_errors=root.parse_errors,
critical_path_tokens=root.critical_path_tokens,
stopped_reason=root.stopped_reason,
)
def _solve_node(self, *, problem: str, depth: int, role: str) -> NodeResult:
if len(self.calls) >= self.max_calls:
return NodeResult(
final_answer=None,
trace="",
critical_path_tokens=0,
parse_errors=[f"max model calls exceeded ({self.max_calls})"],
stopped_reason="max_calls",
)
if depth > self.max_depth:
return NodeResult(
final_answer=None,
trace="",
critical_path_tokens=0,
parse_errors=[f"max recursion depth exceeded ({self.max_depth})"],
stopped_reason="max_depth",
)
input_text = format_model_input(problem, self.prefix_template)
forced_output_prefix = ""
if depth == 0 and role == "root" and self.force_root_action is not None:
forced_output_prefix = FORCED_ROOT_OUTPUT_PREFIXES[self.force_root_action]
output_text = self.generate_once(
input_text,
output_prefix=forced_output_prefix,
)
generated_output = (
output_text[len(forced_output_prefix) :]
if forced_output_prefix and output_text.startswith(forced_output_prefix)
else output_text
)
call = ModelCall(
call_id=len(self.calls) + 1,
depth=depth,
role=role,
problem=problem,
input_text=input_text,
output_text=output_text,
input_tokens=token_count(self.tokenizer, input_text + forced_output_prefix),
output_tokens=token_count(self.tokenizer, generated_output),
forced_output_prefix=forced_output_prefix,
)
self.calls.append(call)
own_tokens = call.input_tokens + call.output_tokens
try:
action_type, paths = parse_model_output(output_text)
validate_action_protocol(output_text, action_type, role)
call.action_type = action_type
except ValueError as exc:
call.parse_error = str(exc)
return NodeResult(
final_answer=None,
trace=output_text,
critical_path_tokens=own_tokens,
parse_errors=[f"call_{call.call_id}: {exc}"],
stopped_reason="parse_error",
)
if action_type in {"SOLVE_DIRECTLY", "AGGREGATE"}:
return NodeResult(
final_answer=extract_final_answer(output_text),
trace=output_text,
critical_path_tokens=own_tokens,
parse_errors=[],
)
child_results = [
self._solve_node(problem=path, depth=depth + 1, role=f"path_{idx}")
for idx, path in enumerate(paths, start=1)
]
aggregate_problem = build_aggregation_problem(problem, paths, child_results)
aggregate_result = self._solve_node(
problem=aggregate_problem,
depth=depth + 1,
role="aggregate",
)
child_trace = "\n\n".join(
f"<problem>\n{path.strip()}\n</problem>\n{result.trace.strip()}"
for path, result in zip(paths, child_results)
)
full_trace = "\n\n".join(
part
for part in [
output_text.strip(),
child_trace.strip(),
f"<problem>\n{aggregate_problem.strip()}\n</problem>\n{aggregate_result.trace.strip()}",
]
if part
)
child_critical = max((result.critical_path_tokens for result in child_results), default=0)
parse_errors: list[str] = []
stopped_reason = "completed"
for result in child_results + [aggregate_result]:
parse_errors.extend(result.parse_errors)
if result.stopped_reason != "completed":
stopped_reason = result.stopped_reason
return NodeResult(
final_answer=aggregate_result.final_answer,
trace=full_trace,
critical_path_tokens=own_tokens + child_critical + aggregate_result.critical_path_tokens,
parse_errors=parse_errors,
stopped_reason=stopped_reason,
)
def load_model_and_tokenizer(model_path: str, dtype: str = "bf16", device_map: str | None = None) -> tuple[Any, Any]:
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
dtype_map = {
"bf16": torch.bfloat16,
"fp16": torch.float16,
"fp32": torch.float32,
"auto": "auto",
}
torch_dtype = dtype_map.get(dtype, torch.bfloat16)
kwargs: dict[str, Any] = {
"torch_dtype": torch_dtype,
"trust_remote_code": True,
}
if device_map:
kwargs["device_map"] = device_map
model = AutoModelForCausalLM.from_pretrained(model_path, **kwargs)
if not device_map:
model = model.to("cuda" if torch.cuda.is_available() else "cpu")
# SFT checkpoints commonly persist use_cache=False from gradient checkpointing.
# Generation should always use the KV cache, or long rollouts become prohibitively slow.
model.config.use_cache = True
model.eval()
return model, tokenizer
def row_to_eval_record(
item: dict[str, Any],
*,
idx: int,
split: str,
problem_field: str | None,
answer_field: str | None,
) -> dict[str, Any]:
p_field = problem_field or infer_field(item, ["problem", "question", "input"])
a_field = answer_field or infer_field(item, ["answer", "target", "final_answer", "solution"])
if p_field is None:
raise ValueError(f"Could not infer problem field. Available fields: {list(item.keys())}")
return {
"id": str(item.get("id", item.get("unique_id", f"{split}_{idx:06d}"))),
"problem": str(item[p_field]),
"answer": str(item[a_field]) if a_field and item.get(a_field) is not None else None,
"raw": dict(item),
}
def load_eval_jsonl(
*,
path: str | Path,
problem_field: str | None,
answer_field: str | None,
limit: int | None,
start_index: int = 0,
) -> list[dict[str, Any]]:
rows: list[dict[str, Any]] = []
path = Path(path)
with path.open("r", encoding="utf-8") as handle:
for absolute_idx, line in enumerate(handle):
if absolute_idx < start_index:
continue
if limit is not None and len(rows) >= limit:
break
if not line.strip():
continue
item = json.loads(line)
rows.append(
row_to_eval_record(
item,
idx=absolute_idx,
split="jsonl",
problem_field=problem_field,
answer_field=answer_field,
)
)
return rows
def load_eval_dataset(
*,
dataset: str,
dataset_config: str | None,
split: str,
problem_field: str | None,
answer_field: str | None,
limit: int | None,
start_index: int = 0,
) -> list[dict[str, Any]]:
from datasets import load_dataset
ds = load_dataset(dataset, dataset_config, split=split) if dataset_config else load_dataset(dataset, split=split)
if start_index:
ds = ds.select(range(start_index, len(ds)))
if limit is not None:
ds = ds.select(range(min(limit, len(ds))))
rows: list[dict[str, Any]] = []
for idx, item in enumerate(ds, start=start_index):
rows.append(
row_to_eval_record(
dict(item),
idx=idx,
split=split,
problem_field=problem_field,
answer_field=answer_field,
)
)
return rows
def infer_field(item: dict[str, Any], candidates: list[str]) -> str | None:
for candidate in candidates:
if candidate in item:
return candidate
return None
def write_jsonl(path: str | Path, rows: list[dict[str, Any]]) -> None:
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
with path.open("w", encoding="utf-8") as handle:
for row in rows:
handle.write(json.dumps(row, ensure_ascii=False) + "\n")