0704

2026-07-04 01:37:11 +08:00
parent 878ad53229
commit b8b0855048
18 changed files with 1598 additions and 1342 deletions
@@ -0,0 +1,178 @@
+#!/usr/bin/env python3
+"""Use local LM Studio OpenAI-compatible API to polish markdown chapters."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import sys
+import time
+import urllib.error
+import urllib.request
+from pathlib import Path
+
+
+DEFAULT_BASE_URL = "http://localhost:1234/v1"
+DEFAULT_MODEL = "gemma4-12b-qat-uncensored-hauhaucs-balanced"
+
+
+def read_text(path: Path) -> str:
+    return path.read_text(encoding="utf-8-sig")
+
+
+def write_text(path: Path, text: str) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(text.rstrip() + "\n", encoding="utf-8")
+
+
+def text_stats(text: str) -> dict[str, int]:
+    return {
+        "chars": len(text),
+        "nonspace": len(re.findall(r"\S", text)),
+        "zh": len(re.findall(r"[\u4e00-\u9fff]", text)),
+    }
+
+
+def request_json(url: str, payload: dict | None = None, timeout: int = 600) -> dict:
+    data = None if payload is None else json.dumps(payload, ensure_ascii=False).encode("utf-8")
+    req = urllib.request.Request(url, data=data)
+    if payload is not None:
+        req.add_header("Content-Type", "application/json")
+    with urllib.request.urlopen(req, timeout=timeout) as resp:
+        raw = resp.read().decode("utf-8")
+    return json.loads(raw)
+
+
+def list_models(base_url: str) -> None:
+    payload = request_json(f"{base_url.rstrip('/')}/models", timeout=10)
+    for item in payload.get("data", []):
+        print(item.get("id", ""))
+
+
+def polish_text(
+    text: str,
+    system_prompt: str,
+    *,
+    base_url: str,
+    model: str,
+    temperature: float,
+    top_p: float | None,
+    presence_penalty: float | None,
+    frequency_penalty: float | None,
+    repeat_penalty: float | None,
+    max_tokens: int | None,
+    timeout: int,
+) -> str:
+    payload = {
+        "model": model,
+        "messages": [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": text},
+        ],
+        "temperature": temperature,
+    }
+    if top_p is not None:
+        payload["top_p"] = top_p
+    if presence_penalty is not None:
+        payload["presence_penalty"] = presence_penalty
+    if frequency_penalty is not None:
+        payload["frequency_penalty"] = frequency_penalty
+    if repeat_penalty is not None:
+        payload["repeat_penalty"] = repeat_penalty
+    if max_tokens is not None:
+        payload["max_tokens"] = max_tokens
+
+    result = request_json(f"{base_url.rstrip('/')}/chat/completions", payload, timeout=timeout)
+    try:
+        return result["choices"][0]["message"]["content"].strip()
+    except (KeyError, IndexError, TypeError) as exc:
+        raise RuntimeError(f"Unexpected LM Studio response: {result}") from exc
+
+
+def collect_inputs(input_path: Path) -> list[Path]:
+    if input_path.is_file():
+        return [input_path]
+    if input_path.is_dir():
+        return sorted(input_path.glob("*.md"))
+    raise FileNotFoundError(f"Input not found: {input_path}")
+
+
+def output_path_for(src: Path, input_root: Path, out: Path) -> Path:
+    if input_root.is_file() and out.suffix.lower() == ".md":
+        return out
+    return out / src.name
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Polish chapters with local LM Studio.")
+    parser.add_argument("--input", "-i", type=Path, help="Input markdown file or directory.")
+    parser.add_argument("--out", "-o", type=Path, default=Path("中间文件/润色输出"), help="Output file or directory.")
+    parser.add_argument("--system", type=Path, default=Path("脚本/lmstudio-润色-system.md"), help="System prompt file.")
+    parser.add_argument("--base-url", default=DEFAULT_BASE_URL)
+    parser.add_argument("--model", default=DEFAULT_MODEL)
+    parser.add_argument("--temperature", type=float, default=1.0)
+    parser.add_argument("--top-p", type=float, default=None)
+    parser.add_argument("--presence-penalty", type=float, default=None)
+    parser.add_argument("--frequency-penalty", type=float, default=None)
+    parser.add_argument("--repeat-penalty", type=float, default=None, help="LM Studio/llama.cpp repeat penalty when supported.")
+    parser.add_argument("--max-tokens", type=int, default=None)
+    parser.add_argument("--timeout", type=int, default=1800)
+    parser.add_argument("--list-models", action="store_true", help="List LM Studio models and exit.")
+    parser.add_argument("--dry-run", action="store_true", help="Show planned tasks without calling LM Studio.")
+    args = parser.parse_args()
+
+    if args.list_models:
+        list_models(args.base_url)
+        return 0
+    if args.input is None:
+        parser.error("--input is required unless --list-models is used")
+
+    inputs = collect_inputs(args.input)
+    if not inputs:
+        raise FileNotFoundError(f"No markdown files under: {args.input}")
+    system_prompt = read_text(args.system)
+    for src in inputs:
+        dst = output_path_for(src, args.input, args.out)
+        print(f"{src} -> {dst}", file=sys.stderr)
+        source_text = read_text(src)
+        source_stats = text_stats(source_text)
+        if args.dry_run:
+            print(
+                f"source chars={source_stats['chars']} nonspace={source_stats['nonspace']} zh={source_stats['zh']}",
+                file=sys.stderr,
+            )
+            continue
+        started = time.time()
+        try:
+            polished = polish_text(
+                source_text,
+                system_prompt,
+                base_url=args.base_url,
+                model=args.model,
+                temperature=args.temperature,
+                top_p=args.top_p,
+                presence_penalty=args.presence_penalty,
+                frequency_penalty=args.frequency_penalty,
+                repeat_penalty=args.repeat_penalty,
+                max_tokens=args.max_tokens,
+                timeout=args.timeout,
+            )
+        except urllib.error.URLError as exc:
+            print(f"LM Studio request failed: {exc}", file=sys.stderr)
+            return 1
+        write_text(dst, polished)
+        polished_stats = text_stats(polished)
+        print(
+            f"wrote {dst} ({time.time() - started:.1f}s) "
+            f"source_zh={source_stats['zh']} output_zh={polished_stats['zh']} "
+            f"source_nonspace={source_stats['nonspace']} output_nonspace={polished_stats['nonspace']}",
+            file=sys.stderr,
+        )
+        print(dst)
+
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())