diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..aa2b704 --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +openai>=1.0.0 diff --git a/scripts/batch_o3_prompts.py b/scripts/batch_o3_prompts.py new file mode 100644 index 0000000..a702102 --- /dev/null +++ b/scripts/batch_o3_prompts.py @@ -0,0 +1,90 @@ +#!/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 + +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:] + 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()