dataset_generate.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
"""Template-based probe dataset generation (no LLM). Writes JSONL to cwd."""
from __future__ import annotations
 
import json
import re
from pathlib import Path
 
_MAX_ROWS = 64
 
_TEMPLATES: list[tuple[str, str]] = [
    ("What is {topic}?", "{topic} is a subject area with core concepts, examples, and practical applications."),
    ("Explain {topic} in simple terms.", "In simple terms, {topic} can be understood as a set of related ideas and practices."),
    ("Give an example related to {topic}.", "A concrete example of {topic} helps illustrate how the concept shows up in practice."),
    ("Why is {topic} important?", "Understanding {topic} supports better decisions, communication, and further learning in the area."),
    ("Describe {topic} to a beginner.", "For a beginner, {topic} is a useful starting point before diving into more advanced material."),
    ("What are key facts about {topic}?", "Key facts about {topic} include definitions, common use cases, and how it connects to nearby topics."),
    ("How does {topic} work?", "How {topic} works depends on context, but it generally follows identifiable patterns and principles."),
    ("List three aspects of {topic}.", "Three aspects of {topic} are foundational ideas, typical examples, and common questions people ask about it."),
]
 
 
def _slug(topic: str) -> str:
    s = re.sub(r"[^\w\s-]", "", topic.lower())
    s = re.sub(r"[\s_-]+", "_", s).strip("_")
    return (s[:48] or "dataset")
 
 
def generate_topic_rows(topic: str, count: int = 5) -> list[dict[str, str]]:
    topic = " ".join(str(topic).split())
    if not topic:
        raise ValueError("topic is required")
    n = max(1, min(int(count), _MAX_ROWS))
    rows: list[dict[str, str]] = []
    for i in range(n):
        instr_tpl, resp_tpl = _TEMPLATES[i % len(_TEMPLATES)]
        rows.append({
            "instruction": instr_tpl.format(topic=topic),
            "response": resp_tpl.format(topic=topic),
        })
    return rows
 
 
def default_output_path(topic: str, cwd: Path | None = None) -> Path:
    base = cwd or Path.cwd()
    return (base / f"{_slug(topic)}_dataset.jsonl").resolve()
 
 
def write_dataset_jsonl(path: Path, rows: list[dict]) -> Path:
    path = path.resolve()
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8", newline="\n") as f:
        for row in rows:
            f.write(json.dumps({"instruction": row["instruction"], "response": row["response"]}, ensure_ascii=False))
            f.write("\n")
    return path
 
 
def run_dataset_generate(
    *,
    topic: str,
    count: int = 5,
    cwd: Path | str | None = None,
    output: str | None = None,
) -> dict:
    rows = generate_topic_rows(topic, count)
    base = Path(cwd).resolve() if cwd else Path.cwd().resolve()
    out_path = (base / output).resolve() if output else default_output_path(topic, base)
    write_dataset_jsonl(out_path, rows)
    return {
        "topic": " ".join(str(topic).split()),
        "count": len(rows),
        "path": str(out_path),
        "rows": rows,
    }