DatasetMaker/engine.py at main · asaad123sh/DatasetMaker · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
import json
import os
import uuid
import random


def _entry_signature(entry):
    signature_payload = {
        "messages": entry.get("messages", []),
        "difficulty": entry.get("difficulty", ""),
        "tags": entry.get("tags", []),
        "safety_level": entry.get("safety_level", "safe"),
        "context": entry.get("context", ""),
    }
    return json.dumps(signature_payload, ensure_ascii=False, sort_keys=True)


def write_category_files(
    base_dir,
    category_name,
    generator_func,
    num_files=100,
    lines_per_file=30000,
    max_unique_attempts=8,
):
    category_dir = os.path.join(base_dir, category_name)
    os.makedirs(category_dir, exist_ok=True)

    total_generated = 0
    for file_idx in range(1, num_files + 1):
        filename = f"{category_name}_{file_idx:04d}.jsonl"
        filepath = os.path.join(category_dir, filename)
        seen_signatures = set()

        with open(filepath, "w", encoding="utf-8") as f:
            for _ in range(lines_per_file):
                entry = None
                signature = None

                for _attempt in range(max_unique_attempts):
                    candidate = generator_func()
                    signature = _entry_signature(candidate)
                    if signature not in seen_signatures:
                        entry = candidate
                        break

                if entry is None:
                    # Fallback keeps generation moving while preserving record uniqueness.
                    candidate = generator_func()
                    base_context = candidate.get("context", "Synthetic template generation")
                    candidate["context"] = f"{base_context} | variation:{uuid.uuid4().hex[:10]}"
                    signature = _entry_signature(candidate)
                    entry = candidate

                seen_signatures.add(signature)

                # Enforce required schema fields
                final_entry = {
                    "id": str(uuid.uuid4()),
                    "source": "synthetic_procedural",
                    "language": "en",
                    "category": category_name.lower(),
                    "messages": entry.get("messages", []),
                    "difficulty": entry.get("difficulty", random.choice(["easy", "medium"])),
                    "quality": "high",
                    "tags": entry.get("tags", []),
                    "safety_level": entry.get("safety_level", "safe"),
                    "context": entry.get("context", "Synthetic template generation")
                }
                f.write(json.dumps(final_entry, ensure_ascii=False) + "\n")
                total_generated += 1

        print(f"[{category_name}] Written {filepath} ({lines_per_file} lines)")

    print(f"Completed {category_name}: {total_generated} lines across {num_files} files.")