from __future__ import annotations import argparse import io import json import sqlite3 import sys from collections.abc import Sequence from contextlib import ExitStack from dataclasses import asdict from pathlib import Path from typing import TextIO from .deberta import MODEL_ID, MODEL_REVISION, DebertaCandidateScorer from .domain import Candidate, CandidateScorer, RerankConfig, RerankRequest, RerankResult from .mozc import MOZC_REVISION, MozcDictionaryIndex, build_mozc_index from .profiles import MOZC_PROFILE_SELECTION_REVISION, MOZC_PROFILES from .reranker import Reranker from .sidecar import serve def _parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(prog="deberta-ime") subparsers = parser.add_subparsers(dest="command", required=True) rerank = subparsers.add_parser("rerank", help="rerank JSON Lines from stdin") rerank.add_argument("--prior-weight", type=float, default=2.0) rerank.add_argument("--min-margin", type=float, default=0.1) rerank.add_argument("--device", default="cpu") rerank.add_argument("--cache-dir", type=Path) rerank.add_argument("--offline", action="store_true") rerank.add_argument("--input", type=Path, help="UTF-8 JSONL file (default: stdin)") rerank.add_argument("--output", type=Path, help="UTF-8 JSONL file (default: stdout)") serve_parser = subparsers.add_parser( "serve", help="run a persistent UTF-8 JSON Lines reranking sidecar" ) serve_parser.add_argument("--device", default="cpu") serve_parser.add_argument("--cache-dir", type=Path) serve_parser.add_argument("--offline", action="store_true") mozc_index = subparsers.add_parser( "mozc-index", help="build a local SQLite index from Mozc OSS dictionary files" ) mozc_index.add_argument("--dictionary-dir", type=Path, required=True) mozc_index.add_argument("--output", type=Path, required=True) mozc_index.add_argument("--source-revision", default=MOZC_REVISION) mozc_rerank = subparsers.add_parser( "mozc-rerank", help="generate finite candidates from a Mozc index and rerank JSON Lines" ) mozc_rerank.add_argument("--index", type=Path, required=True) mozc_rerank.add_argument("--profile", choices=tuple(MOZC_PROFILES), default="incremental") mozc_rerank.add_argument("--candidate-limit", type=int, default=8) mozc_rerank.add_argument("--device", default="cpu") mozc_rerank.add_argument("--cache-dir", type=Path) mozc_rerank.add_argument("--offline", action="store_true") mozc_rerank.add_argument("--input", type=Path, help="UTF-8 JSONL file (default: stdin)") mozc_rerank.add_argument("--output", type=Path, help="UTF-8 JSONL file (default: stdout)") return parser def _request_from_payload(payload: object) -> RerankRequest: if not isinstance(payload, dict): raise ValueError("input must be a JSON object") reading = payload.get("reading") raw_candidates = payload.get("candidates") if not isinstance(reading, str): raise ValueError("reading must be a string") if not isinstance(raw_candidates, list): raise ValueError("candidates must be a list") candidates: list[Candidate] = [] for raw_candidate in raw_candidates: if isinstance(raw_candidate, str): candidates.append(Candidate(raw_candidate)) continue if not isinstance(raw_candidate, dict) or not isinstance( raw_candidate.get("surface"), str ): raise ValueError("each candidate must be a string or an object with surface") prior_score = raw_candidate.get("prior_score", 0.0) if not isinstance(prior_score, int | float): raise ValueError("candidate prior_score must be numeric") candidates.append(Candidate(raw_candidate["surface"], float(prior_score))) left = _string_list(payload.get("left_context", []), "left_context") right = _string_list(payload.get("right_context", []), "right_context") return RerankRequest( reading=reading, candidates=tuple(candidates), left_context=left, right_context=right, ) def _string_list(value: object, name: str) -> tuple[str, ...]: if not isinstance(value, list) or not all(isinstance(item, str) for item in value): raise ValueError(f"{name} must be a list of strings") return tuple(value) def _mozc_request_from_payload( payload: object, index: MozcDictionaryIndex, *, candidate_limit: int, profile: str, ) -> RerankRequest: if not isinstance(payload, dict): raise ValueError("input must be a JSON object") reading = payload.get("reading") if not isinstance(reading, str): raise ValueError("reading must be a string") right_context = _string_list(payload.get("right_context", []), "right_context") if profile == "incremental" and right_context: raise ValueError("incremental profile does not accept right_context") return RerankRequest( reading=reading, candidates=index.lookup(reading, limit=candidate_limit), left_context=_string_list(payload.get("left_context", []), "left_context"), right_context=right_context, ) def _result_payload(result: RerankResult) -> dict[str, object]: return { "ok": True, "schema_version": 1, "model": {"id": MODEL_ID, "revision": MODEL_REVISION}, "reading": result.reading, "decision": { "changed": result.changed, "reason": result.reason, "margin": result.margin, }, "ranked": [ { "surface": item.surface, "original_rank": item.original_rank, "prior_score": item.prior_score, "model_score": item.model_score, "combined_score": item.combined_score, } for item in result.ranked ], } def run( argv: Sequence[str] | None = None, *, stdin: TextIO | None = None, stdout: TextIO | None = None, stderr: TextIO | None = None, scorer: CandidateScorer | None = None, ) -> int: args = _parser().parse_args(argv) error_stream = stderr or sys.stderr if args.command == "serve": active_scorer = scorer or DebertaCandidateScorer( device=args.device, cache_dir=args.cache_dir, local_files_only=args.offline, ) input_stream = stdin if input_stream is None: if isinstance(sys.stdin, io.TextIOWrapper): sys.stdin.reconfigure(encoding="utf-8", errors="strict") input_stream = sys.stdin output_stream = stdout if output_stream is None: if isinstance(sys.stdout, io.TextIOWrapper): sys.stdout.reconfigure( encoding="utf-8", errors="strict", newline="\n", write_through=True, ) output_stream = sys.stdout return serve(input_stream, output_stream, active_scorer) if args.command == "mozc-index": try: manifest = build_mozc_index( args.dictionary_dir, args.output, source_revision=args.source_revision, ) except (OSError, ValueError, sqlite3.Error) as error: error_stream.write(f"error: {error}\n") return 2 output_stream = stdout or sys.stdout output_stream.write( json.dumps( {"ok": True, "schema_version": 1, "manifest": asdict(manifest)}, ensure_ascii=False, separators=(",", ":"), ) ) output_stream.write("\n") return 0 if args.command == "mozc-rerank": if ( args.input is not None and args.output is not None and args.input.resolve() == args.output.resolve() ): error_stream.write("error: --input and --output must be different paths\n") return 2 active_scorer = scorer or DebertaCandidateScorer( device=args.device, cache_dir=args.cache_dir, local_files_only=args.offline, ) selected_profile = MOZC_PROFILES[args.profile] reranker = Reranker(active_scorer, selected_profile) try: index = MozcDictionaryIndex(args.index) except (OSError, ValueError, sqlite3.Error) as error: error_stream.write(f"error: cannot open Mozc index: {error}\n") return 2 with index, ExitStack() as stack: input_stream = ( stack.enter_context(args.input.open(encoding="utf-8")) if args.input is not None else (stdin or sys.stdin) ) output_stream = ( stack.enter_context(args.output.open("w", encoding="utf-8", newline="\n")) if args.output is not None else (stdout or sys.stdout) ) exit_code = 0 for line_number, raw_line in enumerate(input_stream, start=1): if not raw_line.strip(): continue try: request = _mozc_request_from_payload( json.loads(raw_line), index, candidate_limit=args.candidate_limit, profile=args.profile, ) payload = _result_payload(reranker.rerank(request)) payload["candidate_source"] = { "kind": "mozc_oss_dictionary_index", "source_revision": index.manifest.source_revision, "limit": args.candidate_limit, } payload["profile"] = args.profile payload["profile_config"] = { "prior_weight": selected_profile.prior_weight, "min_margin": selected_profile.min_margin, "selected_on_revision": MOZC_PROFILE_SELECTION_REVISION, } except (json.JSONDecodeError, ValueError, TypeError) as error: payload = { "ok": False, "schema_version": 1, "line": line_number, "error": str(error), } exit_code = 2 output_stream.write( json.dumps(payload, ensure_ascii=False, separators=(",", ":")) ) output_stream.write("\n") return exit_code if ( args.input is not None and args.output is not None and args.input.resolve() == args.output.resolve() ): error_stream.write("error: --input and --output must be different paths\n") return 2 active_scorer = scorer or DebertaCandidateScorer( device=args.device, cache_dir=args.cache_dir, local_files_only=args.offline, ) reranker = Reranker( active_scorer, RerankConfig(prior_weight=args.prior_weight, min_margin=args.min_margin), ) with ExitStack() as stack: input_stream = ( stack.enter_context(args.input.open(encoding="utf-8")) if args.input is not None else (stdin or sys.stdin) ) output_stream = ( stack.enter_context(args.output.open("w", encoding="utf-8", newline="\n")) if args.output is not None else (stdout or sys.stdout) ) exit_code = 0 for line_number, raw_line in enumerate(input_stream, start=1): if not raw_line.strip(): continue try: request = _request_from_payload(json.loads(raw_line)) payload = _result_payload(reranker.rerank(request)) except (json.JSONDecodeError, ValueError, TypeError) as error: payload = { "ok": False, "schema_version": 1, "line": line_number, "error": str(error), } exit_code = 2 output_stream.write(json.dumps(payload, ensure_ascii=False, separators=(",", ":"))) output_stream.write("\n") return exit_code def main() -> None: raise SystemExit(run()) if __name__ == "__main__": main()