#!/usr/bin/env python3 """Use local LM Studio OpenAI-compatible API to polish markdown chapters.""" from __future__ import annotations import argparse import difflib 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 generate_patch(src_text: str, polished_text: str, src_path: Path, dst_path: Path) -> str: """Generate a unified diff (patch) between source and polished text.""" src_lines = src_text.splitlines() polished_lines = polished_text.splitlines() diff = difflib.unified_diff( src_lines, polished_lines, fromfile=str(src_path), tofile=str(dst_path), lineterm="", ) return "\n".join(diff) def write_patch(src_text: str, polished_text: str, src_path: Path, dst_path: Path) -> Path: """Write a .patch file next to the polished markdown file.""" patch_path = dst_path.with_suffix(".patch") patch_text = generate_patch(src_text, polished_text, src_path, dst_path) patch_path.write_text(patch_text, encoding="utf-8") return patch_path 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.") parser.add_argument("--no-patch", action="store_true", help="Do not generate .patch files alongside polished output.") 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) if not args.no_patch: patch_path = write_patch(source_text, polished, src, dst) print( f"wrote {dst} + {patch_path} ({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, ) else: 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())