SPB Git

spb/wpai0002_uqo

Public
3.4 KB · 86 lines python
1#!/usr/bin/env python32"""Lightweight local web demo (standard library only): paste an LLM-rewritten passage, get the HTR reconstruction,3a word-level diff and preservation warnings. Run: .venv/bin/python demo/app.py --model <release_dir> --port 8765"""4from __future__ import annotations56import argparse7import difflib8import html9import json10import sys11from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer12from pathlib import Path1314ROOT = Path(__file__).resolve().parents[1]15sys.path.insert(0, str(ROOT / "src"))16from htr.inference.pipeline import NOTICE, HTRReconstructor  # noqa: E4021718PAGE = (Path(__file__).parent / "index.html").read_text()192021def word_diff(a: str, b: str) -> str:22    sm = difflib.SequenceMatcher(None, a.split(), b.split())23    out = []24    for op, i1, i2, j1, j2 in sm.get_opcodes():25        if op == "equal":26            out.append(html.escape(" ".join(a.split()[i1:i2])))27        else:28            if i2 > i1:29                out.append('<del>' + html.escape(" ".join(a.split()[i1:i2])) + "</del>")30            if j2 > j1:31                out.append('<ins>' + html.escape(" ".join(b.split()[j1:j2])) + "</ins>")32    return " ".join(out)333435class H(BaseHTTPRequestHandler):36    rec: HTRReconstructor = None  # type: ignore[assignment]3738    def _send(self, code: int, body: bytes, ctype: str) -> None:39        self.send_response(code)40        self.send_header("Content-Type", ctype)41        self.send_header("Content-Length", str(len(body)))42        self.end_headers()43        self.wfile.write(body)4445    def do_GET(self):  # noqa: N80246        if self.path in ("/", "/index.html"):47            self._send(200, PAGE.replace("{{MODEL}}", html.escape(self.rec.meta.get("name", "htr"))).encode(), "text/html; charset=utf-8")48        elif self.path == "/health":49            self._send(200, json.dumps({"ok": True, "model": self.rec.meta.get("name")}).encode(), "application/json")50        else:51            self._send(404, b"not found", "text/plain")5253    def do_POST(self):  # noqa: N80254        if self.path != "/api/reconstruct":55            return self._send(404, b"not found", "text/plain")56        n = int(self.headers.get("Content-Length", "0"))57        req = json.loads(self.rfile.read(n) or b"{}")58        text = (req.get("text") or "").strip()59        if not text:60            return self._send(400, json.dumps({"error": "empty text"}).encode(), "application/json")61        meta = {k: req.get(k) for k in ("discipline", "pub_year", "section_role") if req.get(k)}62        if "pub_year" in meta:63            meta["pub_year"] = int(meta["pub_year"])64        res = self.rec.reconstruct(text[:6000], meta, beam=int(req.get("beam", 1)))65        body = {**res.to_dict(), "diff_html": word_diff(text, res.reconstruction), "notice": NOTICE}66        self._send(200, json.dumps(body, ensure_ascii=False).encode(), "application/json")6768    def log_message(self, *a):  # quiet69        pass707172def main() -> None:73    ap = argparse.ArgumentParser()74    ap.add_argument("--model", required=True)75    ap.add_argument("--port", type=int, default=8765)76    ap.add_argument("--host", default="127.0.0.1")77    ap.add_argument("--device", default="auto")78    a = ap.parse_args()79    H.rec = HTRReconstructor.from_pretrained(a.model, device=a.device)80    print(f"HTR demo on http://{a.host}:{a.port}  (model {H.rec.meta.get('name')})")81    ThreadingHTTPServer((a.host, a.port), H).serve_forever()828384if __name__ == "__main__":85    main()86
Files · main
Menu

Theme