| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Minimal reference server for watermark detection.""" |
| 5 | |
| 6 | import argparse |
| 7 | |
| 8 | import uvicorn |
| 9 | from fastapi import FastAPI |
| 10 | from pydantic import BaseModel |
| 11 | |
| 12 | from vllm.tokenizers import TokenizerLike, cached_get_tokenizer |
| 13 | from vllm.v1.watermarking import GumbelWatermarkDetector |
| 14 | |
| 15 | app = FastAPI() |
| 16 | tokenizer: TokenizerLike | None = None |
| 17 | detector: GumbelWatermarkDetector | None = None |
| 18 | |
| 19 | |
| 20 | class DetectionRequest(BaseModel): |
| 21 | text: str |
| 22 | |
| 23 | |
| 24 | class DetectionResponse(BaseModel): |
| 25 | score: float |
| 26 | p_value: float |
| 27 | num_scored_tokens: int |
| 28 | is_watermarked: bool |
| 29 | |
| 30 | |
| 31 | @app.post("/detect") |
| 32 | def detect(request: DetectionRequest) -> DetectionResponse: |
| 33 | assert tokenizer is not None |
| 34 | assert detector is not None |
| 35 | token_ids = tokenizer.encode(request.text, add_special_tokens=False) |
| 36 | result = detector.detect(token_ids) |
| 37 | return DetectionResponse( |
| 38 | score=result.score, |
| 39 | p_value=result.p_value, |
| 40 | num_scored_tokens=result.num_scored_tokens, |
| 41 | is_watermarked=result.is_watermarked, |
| 42 | ) |
| 43 | |
| 44 | |
| 45 | def parse_args() -> argparse.Namespace: |
| 46 | parser = argparse.ArgumentParser(description=__doc__) |
| 47 | parser.add_argument("--tokenizer", required=True) |
| 48 | parser.add_argument("--key", required=True, type=int) |
| 49 | parser.add_argument("--prf", choices=("philox",), default="philox") |
| 50 | parser.add_argument("--context-width", type=int, default=4) |
| 51 | parser.add_argument("--p-value-threshold", type=float, default=0.01) |
| 52 | parser.add_argument("--host", default="127.0.0.1") |
| 53 | parser.add_argument("--port", type=int, default=8000) |
| 54 | return parser.parse_args() |
| 55 | |
| 56 | |
| 57 | def main(args: argparse.Namespace) -> None: |
| 58 | global tokenizer, detector |
| 59 | tokenizer = cached_get_tokenizer(args.tokenizer) |
| 60 | detector = GumbelWatermarkDetector( |
| 61 | key=args.key, |
| 62 | context_width=args.context_width, |
| 63 | p_value_threshold=args.p_value_threshold, |
| 64 | prf=args.prf, |
| 65 | ) |
| 66 | uvicorn.run(app, host=args.host, port=args.port) |
| 67 | |
| 68 | |
| 69 | if __name__ == "__main__": |
| 70 | main(parse_args()) |