PLSQL-Unit-Test-Dataset-Maker/main.py at main · BasuruK/PLSQL-Unit-Test-Dataset-Maker · 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
import os
from utils.file_io import get_settings
from utils.datasets import combine_artifacts, create_reasoning_dataset, load_and_format_conversation_dataset, load_and_format_reasoning_cot_dataset, load_and_format_openai_messages_dataset
from utils.thread_worker import scan_for_unit_tests

def main():
    (
        extract_methods,
        make_dataset,
        make_reasoning_dataset,
        make_conversation_dataset,
        make_reasoning_cot_dataset,
        make_open_ai_dataset,
        workspace_dir,
    ) = get_settings()

    if extract_methods:
        scan_for_unit_tests(workspace_dir)
    if make_dataset:
        # Assemble all the generated JSON files into one
        print("\033[97;104m" + "Extraction completed, preparing the Dataset" + "\033[0m")
        combine_artifacts()
    if make_reasoning_dataset:
        # Create reasoning dataset from the combined dataset
        print("\033[97;104m" + "Making Reasoning Dataset from combined dataset" + "\033[0m")
        # You can adjust the number of threads based on your system and API rate limits
        create_reasoning_dataset("dataset/combined_tests.jsonl", "dataset/reasoning_dataset.jsonl", num_threads=80)
    if make_conversation_dataset:
        # Create conversation dataset from the reasoning dataset
        print("\033[97;104m" + "Making Conversation Dataset from reasoning dataset" + "\033[0m")
        load_and_format_conversation_dataset("dataset/reasoning_dataset.jsonl", "dataset/conversation_dataset.jsonl")
    # Also convert to OpenAI messages-only format if requested by setting (handled below)
    if make_reasoning_cot_dataset:
        # Create reasoning chain-of-thought dataset from the reasoning dataset
        print("\033[97;104m" + "Making Reasoning CoT Dataset from reasoning dataset" + "\033[0m")
        load_and_format_reasoning_cot_dataset("dataset/reasoning_dataset.jsonl", "dataset/reasoning_cot_dataset.jsonl")
    if make_open_ai_dataset:
        # Create OpenAI messages dataset from the conversation dataset
        print("\033[97;104m" + "Making OpenAI Messages Dataset from conversation dataset" + "\033[0m")
        load_and_format_openai_messages_dataset("dataset/conversation_dataset.jsonl", "dataset/openai_messages_dataset.jsonl", enable_token_check=True)


if __name__ == "__main__":
    # Main function acts as the main entry point to the program
    main()