#!/usr/bin/env python3
"""Run PertMind with vLLM."""

from __future__ import annotations

import argparse

from vllm import LLM, SamplingParams


DEFAULT_SYSTEM_PROMPT = (
    "You are PertMind, a biomedical assistant. For biomedical prediction, "
    "screen-ranking, or gene-set interpretation tasks, answer first and then "
    "provide a concise explanation. Use this style when applicable:\n"
    "Final Answer: <answer>\nExplanation: <brief explanation>"
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", default=".", help="Path or Hugging Face model id.")
    parser.add_argument("--prompt", required=True, help="User prompt.")
    parser.add_argument("--system-prompt", default=DEFAULT_SYSTEM_PROMPT)
    parser.add_argument("--max-tokens", type=int, default=768)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--top-p", type=float, default=0.95)
    parser.add_argument("--max-model-len", type=int, default=12288)
    parser.add_argument("--gpu-memory-utilization", type=float, default=0.85)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    llm = LLM(
        model=args.model,
        trust_remote_code=True,
        dtype="bfloat16",
        max_model_len=args.max_model_len,
        gpu_memory_utilization=args.gpu_memory_utilization,
    )
    tokenizer = llm.get_tokenizer()
    messages = [
        {"role": "system", "content": args.system_prompt},
        {"role": "user", "content": args.prompt},
    ]
    try:
        text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
            enable_thinking=False,
        )
    except TypeError:
        text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    params = SamplingParams(
        temperature=args.temperature,
        top_p=args.top_p,
        max_tokens=args.max_tokens,
    )
    output = llm.generate([text], params)[0].outputs[0].text.strip()
    print(output)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())