File size: 1,480 Bytes
c29257f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | #!/usr/bin/env python3
"""Run one AREX BrowseComp generation through an OpenAI-compatible endpoint."""
from __future__ import annotations
import argparse
import os
from prompts import build_messages
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--question", required=True)
parser.add_argument(
"--base-url",
default=os.environ.get("AREX_BASE_URL", "http://127.0.0.1:8000/v1"),
)
parser.add_argument(
"--api-key",
default=os.environ.get("AREX_API_KEY", "EMPTY"),
)
parser.add_argument(
"--model",
default=os.environ.get("AREX_MODEL", "AREX-Base"),
)
parser.add_argument("--max-tokens", type=int, default=8192)
return parser.parse_args()
def main() -> None:
args = parse_args()
try:
from openai import OpenAI
except ImportError as exc:
raise SystemExit("Install the client first: pip install -U openai") from exc
client = OpenAI(
base_url=args.base_url.rstrip("/") + "/",
api_key=args.api_key,
timeout=600.0,
)
response = client.chat.completions.create(
model=args.model,
messages=build_messages(args.question),
max_tokens=args.max_tokens,
temperature=1.0,
top_p=0.95,
presence_penalty=1.5,
extra_body={"top_k": 20},
)
print(response.choices[0].message.content or "")
if __name__ == "__main__":
main()
|