models/Bidirectional-Parallel-Refresh/inference.py

#!/usr/bin/env python3
"""Call the local Bidirectional-Parallel-Refresh vLLM API."""

import argparse
import json
from pathlib import Path
from typing import Any

try:
    from openai import OpenAI
except ImportError as exc:
    raise SystemExit(
        "The openai package is required. Run with /workspace/vllm/.venv/bin/python "
        "inside the ms-swift container."
    ) from exc


SCRIPT_DIR = Path(__file__).resolve().parent
DEFAULT_PROMPT = (
    "Let x and y be real numbers satisfying x + y = 10 and xy = 21. "
    "Compute x^2 + y^2 and explain your reasoning."
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Generate text from a Bidirectional-Parallel-Refresh "
            "vLLM OpenAI-style endpoint."
        )
    )
    parser.add_argument(
        "--port",
        type=int,
        required=True,
        help="Port printed by start_vllm_server.sh.",
    )
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--model", default="dmtd-qwen3-4b")
    parser.add_argument("--max-gen-tokens", type=int, default=256)
    parser.add_argument("--prompt", default=DEFAULT_PROMPT)
    parser.add_argument(
        "--system-prompt",
        default="You are a careful mathematical reasoning assistant.",
    )
    parser.add_argument("--api-key", default="EMPTY")
    parser.add_argument("--timeout", type=float, default=300.0)
    parser.add_argument("--seed", type=int)
    parser.add_argument(
        "--generation-config",
        type=Path,
        default=SCRIPT_DIR / "generation_config.json",
    )
    return parser.parse_args()


def load_generation_config(path: Path) -> dict[str, Any]:
    try:
        with path.open(encoding="utf-8") as config_file:
            return json.load(config_file)
    except (OSError, json.JSONDecodeError) as exc:
        raise SystemExit(f"Could not read generation config {path}: {exc}") from exc


def main() -> None:
    args = parse_args()
    if not 1 <= args.port <= 65535:
        raise SystemExit(f"Invalid port: {args.port}")
    if args.max_gen_tokens < 1:
        raise SystemExit("--max-gen-tokens must be positive")

    generation_config = load_generation_config(args.generation_config)
    do_sample = bool(generation_config.get("do_sample", False))
    temperature = float(generation_config.get("temperature", 1.0)) if do_sample else 0
    top_p = float(generation_config.get("top_p", 1.0)) if do_sample else 1.0

    extra_body: dict[str, Any] = {}
    if do_sample and "top_k" in generation_config:
        extra_body["top_k"] = int(generation_config["top_k"])
    if "eos_token_id" in generation_config:
        eos_token_ids = generation_config["eos_token_id"]
        if isinstance(eos_token_ids, int):
            eos_token_ids = [eos_token_ids]
        extra_body["stop_token_ids"] = eos_token_ids
    if args.seed is not None:
        extra_body["seed"] = args.seed

    endpoint = f"http://{args.host}:{args.port}/v1"
    client = OpenAI(
        api_key=args.api_key,
        base_url=endpoint,
        timeout=args.timeout,
    )
    response = client.chat.completions.create(
        model=args.model,
        messages=[
            {"role": "system", "content": args.system_prompt},
            {"role": "user", "content": args.prompt},
        ],
        max_tokens=args.max_gen_tokens,
        temperature=temperature,
        top_p=top_p,
        extra_body=extra_body,
    )

    print(f"Endpoint: {endpoint}")
    print(
        "Generation config: "
        f"do_sample={do_sample}, temperature={temperature}, top_p={top_p}, "
        f"top_k={extra_body.get('top_k')}, "
        f"eos_token_id={extra_body.get('stop_token_ids')}"
    )
    print("\nResponse:\n")
    print(response.choices[0].message.content or "")
    if response.usage is not None:
        print(
            "\nUsage: "
            f"prompt={response.usage.prompt_tokens}, "
            f"completion={response.usage.completion_tokens}, "
            f"total={response.usage.total_tokens}"
        )


if __name__ == "__main__":
    main()