generate_predictions_example.py

"""Generate predictions from a vLLM model for the Chinese Writing Benchmark."""

import json
import asyncio
from openai import AsyncOpenAI
from datasets import load_dataset
from tqdm import tqdm

API_BASE = "YOUR_API_BASE"
API_KEY = "YOUR_API_KEY"
MODEL_NAME = "YOUR_MODEL_NAME"
MAX_CONCURRENCY = 32
OUTPUT_PATH = "predictions.json"


async def main():
    ds = load_dataset("zake7749/chinese-writing-benchmark", split="gpt_4.1_2025_04_14")
    prompts = ds["prompt"]
    print(f"Loaded {len(prompts)} prompts from the benchmark dataset.")

    client = AsyncOpenAI(base_url=API_BASE, api_key=API_KEY)
    semaphore = asyncio.Semaphore(MAX_CONCURRENCY)
    prompt_to_response = {}

    async def infer(prompt: str) -> tuple[str, str]:
        async with semaphore:
            resp = await client.chat.completions.create(
                model=MODEL_NAME,
                messages=[{"role": "user", "content": prompt}],
                temperature=0.7,
                max_tokens=16384,
            )
            return prompt, resp.choices[0].message.content

    tasks = [asyncio.create_task(infer(p)) for p in prompts]

    for fut in tqdm(asyncio.as_completed(tasks), total=len(tasks), desc="Generating"):
        prompt, response = await fut
        prompt_to_response[prompt] = response

    with open(OUTPUT_PATH, "w", encoding="utf-8") as f:
        json.dump(prompt_to_response, f, ensure_ascii=False, indent=2)

    print(f"Done! Saved {len(prompt_to_response)} predictions to {OUTPUT_PATH}")


if __name__ == "__main__":
    asyncio.run(main())