258 lines
10 KiB
Python
258 lines
10 KiB
Python
"""数据处理任务的大模型生成适配器。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import json
|
||
import re
|
||
from collections.abc import Callable, Iterable, Mapping
|
||
from typing import Any
|
||
from urllib.parse import urlsplit, urlunsplit
|
||
|
||
import httpx
|
||
|
||
from app.modules.data_process.algorithms import normalize_text, stable_split_assignments
|
||
|
||
|
||
class ModelGenerationError(ValueError):
|
||
"""模型配置、响应或调用失败。"""
|
||
|
||
|
||
def chat_completions_url(value: str) -> str:
|
||
"""把域名、基础 URL 或完整地址统一为 chat completions 地址。"""
|
||
|
||
raw = (value or "").strip()
|
||
if not raw:
|
||
raise ModelGenerationError("generation model api_url is required")
|
||
if "://" not in raw:
|
||
raw = f"https://{raw}"
|
||
parsed = urlsplit(raw)
|
||
if parsed.scheme not in {"http", "https"} or not parsed.hostname:
|
||
raise ModelGenerationError("generation model api_url must be an HTTP(S) host or URL")
|
||
if parsed.username or parsed.password:
|
||
raise ModelGenerationError("generation model api_url must not contain credentials")
|
||
|
||
path = parsed.path.rstrip("/")
|
||
if path.endswith("/chat/completions"):
|
||
target_path = path
|
||
elif path.endswith("/v1"):
|
||
target_path = f"{path}/chat/completions"
|
||
elif not path:
|
||
target_path = "/v1/chat/completions"
|
||
else:
|
||
target_path = f"{path}/v1/chat/completions"
|
||
return urlunsplit((parsed.scheme, parsed.netloc, target_path, "", ""))
|
||
|
||
|
||
def _message_content(payload: Mapping[str, Any]) -> str:
|
||
try:
|
||
content = payload["choices"][0]["message"]["content"]
|
||
except (KeyError, IndexError, TypeError) as exc:
|
||
raise ModelGenerationError("model response does not contain choices[0].message.content") from exc
|
||
if isinstance(content, str):
|
||
return content
|
||
if isinstance(content, list):
|
||
parts = [
|
||
str(item.get("text") or "")
|
||
for item in content
|
||
if isinstance(item, Mapping) and item.get("type") in {None, "text", "output_text"}
|
||
]
|
||
if parts:
|
||
return "".join(parts)
|
||
raise ModelGenerationError("model response content must be text")
|
||
|
||
|
||
def _json_payload(content: str) -> Any:
|
||
cleaned = re.sub(r"<think>[\s\S]*?</think>", "", content, flags=re.IGNORECASE).strip()
|
||
fenced = re.fullmatch(r"```(?:json)?\s*([\s\S]*?)\s*```", cleaned, flags=re.IGNORECASE)
|
||
if fenced:
|
||
cleaned = fenced.group(1).strip()
|
||
try:
|
||
return json.loads(cleaned)
|
||
except json.JSONDecodeError as exc:
|
||
raise ModelGenerationError(
|
||
f"model response is not valid JSON at line {exc.lineno}, column {exc.colno}"
|
||
) from exc
|
||
|
||
|
||
def _result_items(payload: Any) -> list[Mapping[str, Any]]:
|
||
if isinstance(payload, list):
|
||
values = payload
|
||
elif isinstance(payload, Mapping):
|
||
nested = next(
|
||
(
|
||
payload[key]
|
||
for key in ("items", "results", "data", "records")
|
||
if isinstance(payload.get(key), list)
|
||
),
|
||
None,
|
||
)
|
||
values = nested if isinstance(nested, list) else [payload]
|
||
else:
|
||
raise ModelGenerationError("model JSON must be an object or array")
|
||
items = [item for item in values if isinstance(item, Mapping)]
|
||
if not items:
|
||
raise ModelGenerationError("model JSON does not contain result objects")
|
||
return items
|
||
|
||
|
||
def _prompt_messages(prompt: str, content: str, count: int) -> list[dict[str, str]]:
|
||
schema_instruction = (
|
||
f"必须只返回 JSON 对象，格式为 {{\"items\":[{{\"instruction\":\"...\","
|
||
f"\"input\":\"...\",\"output\":\"...\"}}]}}；items 必须包含 {count} 条。"
|
||
"instruction 和 output 不得为空，不要输出 Markdown 代码围栏或分析过程。"
|
||
)
|
||
base_prompt = (
|
||
normalize_text(prompt)
|
||
or "请根据来源内容生成可用于监督微调的问答数据。"
|
||
)
|
||
if "{{ content }}" in base_prompt:
|
||
user_prompt = base_prompt.replace("{{ content }}", content)
|
||
return [
|
||
{"role": "system", "content": schema_instruction},
|
||
{"role": "user", "content": user_prompt},
|
||
]
|
||
return [
|
||
{"role": "system", "content": f"{base_prompt}\n{schema_instruction}"},
|
||
{"role": "user", "content": f"来源内容：\n{content}"},
|
||
]
|
||
|
||
|
||
def generate_model_records(
|
||
preview_items: Iterable[Mapping[str, Any]],
|
||
*,
|
||
model: Mapping[str, Any],
|
||
config: Mapping[str, Any],
|
||
task_id: str,
|
||
split: Mapping[str, int],
|
||
qa_pairs_per_item: int,
|
||
client: httpx.Client | None = None,
|
||
on_progress: Callable[[int, int], None] | None = None,
|
||
) -> list[dict[str, Any]]:
|
||
"""调用 OpenAI 兼容接口，将预览切片生成标准训练记录。
|
||
|
||
单条调用失败会产生可人工修复的 invalid 结果，不会丢弃整批任务。
|
||
"""
|
||
|
||
if not 1 <= qa_pairs_per_item <= 5:
|
||
raise ModelGenerationError("qa_pairs_per_item must be in [1, 5]")
|
||
endpoint = chat_completions_url(str(model.get("api_url") or ""))
|
||
model_name = str(model.get("online_model_name") or model.get("name") or "").strip()
|
||
if not model_name:
|
||
raise ModelGenerationError("generation model name is required")
|
||
|
||
temperature = float(config.get("temperature", 0.7))
|
||
max_tokens = int(config.get("max_tokens", 1024))
|
||
timeout = max(1.0, min(120.0, float(config.get("request_timeout_seconds", 60))))
|
||
retries = max(0, min(5, int(config.get("generation_retries", 2))))
|
||
headers = {"Content-Type": "application/json"}
|
||
api_key = str(model.get("api_key") or "").strip()
|
||
if api_key:
|
||
headers["Authorization"] = f"Bearer {api_key}"
|
||
|
||
owns_client = client is None
|
||
http_client = client or httpx.Client(timeout=timeout)
|
||
results: list[dict[str, Any]] = []
|
||
try:
|
||
preview_list = list(preview_items)
|
||
total_items = len(preview_list)
|
||
for item_index, item in enumerate(preview_list):
|
||
preview_id = str(item.get("id") or f"preview-{item_index + 1}")
|
||
content = normalize_text(
|
||
str(item.get("edited_content") or item.get("original_content") or "")
|
||
)
|
||
request_payload: dict[str, Any] = {
|
||
"model": model_name,
|
||
"messages": _prompt_messages(
|
||
str(config.get("generation_prompt") or ""),
|
||
content,
|
||
qa_pairs_per_item,
|
||
),
|
||
"temperature": temperature,
|
||
"max_tokens": max_tokens,
|
||
}
|
||
if bool(config.get("json_mode", False)):
|
||
request_payload["response_format"] = {"type": "json_object"}
|
||
|
||
last_error: Exception | None = None
|
||
generated_items: list[Mapping[str, Any]] | None = None
|
||
for _ in range(retries + 1):
|
||
try:
|
||
response = http_client.post(endpoint, headers=headers, json=request_payload)
|
||
response.raise_for_status()
|
||
body = response.json()
|
||
if not isinstance(body, Mapping):
|
||
raise ModelGenerationError("model response body must be a JSON object")
|
||
generated_items = _result_items(_json_payload(_message_content(body)))
|
||
break
|
||
except (httpx.HTTPError, json.JSONDecodeError, ModelGenerationError) as exc:
|
||
last_error = exc
|
||
|
||
if generated_items is None:
|
||
error_message = str(last_error or "model generation failed")[:2000]
|
||
result_id = f"result_{hashlib.sha256(f'{preview_id}:error'.encode()).hexdigest()[:16]}"
|
||
results.append(
|
||
{
|
||
"id": result_id,
|
||
"preview_item_id": preview_id,
|
||
"instruction": "模型生成失败，请人工补充",
|
||
"input": content,
|
||
"output": "",
|
||
"original_instruction": "模型生成失败，请人工补充",
|
||
"original_input": content,
|
||
"original_output": "",
|
||
"status": "invalid",
|
||
"error": error_message,
|
||
"split": "train",
|
||
}
|
||
)
|
||
if on_progress:
|
||
on_progress(item_index + 1, total_items)
|
||
continue
|
||
|
||
for variant_index, value in enumerate(generated_items[:qa_pairs_per_item]):
|
||
instruction = normalize_text(str(value.get("instruction") or value.get("question") or ""))
|
||
input_text = normalize_text(str(value.get("input") or value.get("context") or ""))
|
||
output = normalize_text(
|
||
str(
|
||
value.get("output")
|
||
or value.get("answer")
|
||
or value.get("response")
|
||
or ""
|
||
)
|
||
)
|
||
raw_id = f"{preview_id}:{variant_index + 1}:{instruction}:{output}"
|
||
result_id = f"result_{hashlib.sha256(raw_id.encode()).hexdigest()[:16]}"
|
||
valid = bool(instruction and output)
|
||
results.append(
|
||
{
|
||
"id": result_id,
|
||
"preview_item_id": preview_id,
|
||
"instruction": instruction,
|
||
"input": input_text,
|
||
"output": output,
|
||
"original_instruction": instruction,
|
||
"original_input": input_text,
|
||
"original_output": output,
|
||
"status": "valid" if valid else "invalid",
|
||
"error": None if valid else "model result is missing instruction or output",
|
||
"split": "train",
|
||
}
|
||
)
|
||
if on_progress:
|
||
on_progress(item_index + 1, total_items)
|
||
finally:
|
||
if owns_client:
|
||
http_client.close()
|
||
assignments = stable_split_assignments(
|
||
[str(result["id"]) for result in results],
|
||
split,
|
||
seed=task_id,
|
||
)
|
||
for result, assignment in zip(results, assignments, strict=True):
|
||
result["split"] = assignment
|
||
return results
|
||
|
||
|
||
__all__ = ["ModelGenerationError", "chat_completions_url", "generate_model_records"]
|