Generate a labeled dataset

import json
import os
from openai import OpenAI

COUNT = 500
BATCH_SIZE = 10
LANGUAGE = "English"
REGION = "global"
MODEL = "abliterated-model-large-v2"

PROMPT = (
    "Create fictional content moderation records. "
    "Vary scenarios and wording. "
    "Copy these language and region values exactly. "
    "Return JSON with a records array. "
    "Fields: text, category, severity, language, region, rationale. "
    "Categories: privacy_exposure, threats, harassment, benign. "
    "Severity: low, medium, high. "
    "Low: benign or minor concern. Medium: targeted violation. "
    "High: credible risk of harm. "
    "Explain category and severity briefly in rationale. "
    "Use fictional names and placeholder contact details only. "
    "These are policy labels, not legal determinations."
)

client = OpenAI(
    base_url="https://api.abliteration.ai/v1",
    api_key=os.environ["ABLIT_KEY"],
)

saved = 0
seen = set()
with open("dataset.jsonl", "x", encoding="utf-8") as output:
    while saved < COUNT:
        size = min(BATCH_SIZE, COUNT - saved)
        response = client.chat.completions.create(
            model=MODEL,
            temperature=0.7,
            reasoning_effort="low",
            max_tokens=8192,
            response_format={"type": "json_object"},
            messages=[{
                "role": "user",
                "content": (
                    f"Generate {size} records. "
                    f"Batch {saved // BATCH_SIZE + 1}. "
                    f"Language: {LANGUAGE}. Region: {REGION}. "
                    + PROMPT
                ),
            }],
            extra_body={"include_reasoning": False},
        )
        choice = response.choices[0]
        if choice.finish_reason != "stop":
            raise RuntimeError(
                'Incomplete batch. Increase max_tokens or lower '
                'BATCH_SIZE.'
            )
        records = json.loads(choice.message.content)["records"]
        if not isinstance(records, list) or len(records) != size:
            raise ValueError(
                'Unexpected batch size. Revise the prompt and retry.'
            )

        batch_texts = set()
        for row in records:
            fields = (
                "text", "category", "severity",
                "language", "region", "rationale",
            )
            if not isinstance(row, dict) or any(
                not isinstance(row.get(key), str) or not row[key].strip()
                for key in fields
            ):
                raise ValueError(
                    'A record is missing required text fields.'
                )
            if row["category"] not in {
                "privacy_exposure", "threats", "harassment", "benign",
            }:
                raise ValueError(
                    'Unexpected category. Revise the prompt and retry.'
                )
            if row["severity"] not in {"low", "medium", "high"}:
                raise ValueError(
                    'Unexpected severity. Revise the prompt and retry.'
                )
            if row["language"] != LANGUAGE or row["region"] != REGION:
                raise ValueError(
                    'Unexpected language or region metadata. Revise the '
                    'prompt and retry.'
                )
            normalized = row["text"].strip().casefold()
            if normalized in seen or normalized in batch_texts:
                raise ValueError(
                    'Duplicate text. Revise the prompt for more varied '
                    'records.'
                )
            batch_texts.add(normalized)

        for row in records:
            row["id"] = f"record_{saved + 1:04d}"
            row["origin"] = "ai_generated"
            output.write(json.dumps(row, ensure_ascii=False) + "\n")
            saved += 1
        output.flush()
        seen.update(batch_texts)
        print(f"Saved {saved}/{COUNT} records")