93 lines
2.5 KiB
Python
93 lines
2.5 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Batch runner for OpenAI o3 model prompts.
|
||
|
||
Usage:
|
||
python scripts/batch_o3_prompts.py template.txt inputs.jsonl outputs.jsonl
|
||
|
||
• template.txt – prompt template using str.format placeholders
|
||
• inputs.jsonl – JSON-Lines; each line supplies values for the template
|
||
• outputs.jsonl – appended with one JSON object per request (input, prompt, response)
|
||
|
||
Requires the OPENAI_API_KEY environment variable.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
from typing import Any, Dict
|
||
|
||
from dotenv import load_dotenv
|
||
import openai
|
||
|
||
MODEL = "o3"
|
||
|
||
|
||
def load_template(path: Path) -> str:
|
||
with path.open(encoding="utf-8") as fh:
|
||
return fh.read()
|
||
|
||
|
||
def stream_inputs(path: Path):
|
||
with path.open(encoding="utf-8") as fh:
|
||
for line in fh:
|
||
if line.strip():
|
||
yield json.loads(line)
|
||
|
||
|
||
def write_output(path: Path, record: Dict[str, Any]) -> None:
|
||
with path.open("a", encoding="utf-8") as fh:
|
||
fh.write(json.dumps(record, ensure_ascii=False) + "\n")
|
||
|
||
|
||
def call_o3(prompt: str) -> str:
|
||
response = openai.Completion.create(
|
||
model=MODEL,
|
||
prompt=prompt,
|
||
max_tokens=1024,
|
||
temperature=0.7,
|
||
)
|
||
return response.choices[0].text.strip()
|
||
|
||
|
||
def main(argv: list[str] | None = None) -> None:
|
||
argv = argv or sys.argv[1:]
|
||
load_dotenv() # read variables from .env if present
|
||
if len(argv) != 3:
|
||
print(main.__doc__)
|
||
sys.exit(1)
|
||
|
||
template_path, inputs_path, outputs_path = map(Path, argv)
|
||
|
||
if "OPENAI_API_KEY" not in os.environ:
|
||
print("Error: OPENAI_API_KEY environment variable is not set.", file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
template = load_template(template_path)
|
||
|
||
for input_obj in stream_inputs(inputs_path):
|
||
prompt = template.format(**input_obj)
|
||
try:
|
||
answer = call_o3(prompt)
|
||
write_output(
|
||
outputs_path,
|
||
{"input": input_obj, "prompt": prompt, "response": answer},
|
||
)
|
||
print(f"Processed: {input_obj}", file=sys.stderr)
|
||
except Exception as exc: # noqa: BLE001
|
||
write_output(
|
||
outputs_path,
|
||
{
|
||
"input": input_obj,
|
||
"prompt": prompt,
|
||
"error": str(exc),
|
||
},
|
||
)
|
||
print(f"Error processing {input_obj}: {exc}", file=sys.stderr)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|