aurora

1
import os
2
import json
3
import gradio as gr
4
from typing import TYPE_CHECKING, Any, Dict, Tuple
5

6
from llmtuner.webui.common import DATA_CONFIG
7

8
if TYPE_CHECKING:
9
    from gradio.components import Component
10

11

12
PAGE_SIZE = 2
13

14

15
def prev_page(page_index: int) -> int:
16
    return page_index - 1 if page_index > 0 else page_index
17

18

19
def next_page(page_index: int, total_num: int) -> int:
20
    return page_index + 1 if (page_index + 1) * PAGE_SIZE < total_num else page_index
21

22

23
def can_preview(dataset_dir: str, dataset: list) -> Dict[str, Any]:
24
    with open(os.path.join(dataset_dir, DATA_CONFIG), "r", encoding="utf-8") as f:
25
        dataset_info = json.load(f)
26

27
    if (
28
        len(dataset) > 0
29
        and "file_name" in dataset_info[dataset[0]]
30
        and os.path.isfile(os.path.join(dataset_dir, dataset_info[dataset[0]]["file_name"]))
31
    ):
32
        return gr.update(interactive=True)
33
    else:
34
        return gr.update(interactive=False)
35

36

37
def get_preview(dataset_dir: str, dataset: list, page_index: int) -> Tuple[int, list, Dict[str, Any]]:
38
    with open(os.path.join(dataset_dir, DATA_CONFIG), "r", encoding="utf-8") as f:
39
        dataset_info = json.load(f)
40

41
    data_file: str = dataset_info[dataset[0]]["file_name"]
42
    with open(os.path.join(dataset_dir, data_file), "r", encoding="utf-8") as f:
43
        if data_file.endswith(".json"):
44
            data = json.load(f)
45
        elif data_file.endswith(".jsonl"):
46
            data = [json.loads(line) for line in f]
47
        else:
48
            data = [line for line in f]
49
    return len(data), data[PAGE_SIZE * page_index : PAGE_SIZE * (page_index + 1)], gr.update(visible=True)
50

51

52
def create_preview_box(dataset_dir: "gr.Textbox", dataset: "gr.Dropdown") -> Dict[str, "Component"]:
53
    data_preview_btn = gr.Button(interactive=False, scale=1)
54
    with gr.Column(visible=False, elem_classes="modal-box") as preview_box:
55
        with gr.Row():
56
            preview_count = gr.Number(value=0, interactive=False, precision=0)
57
            page_index = gr.Number(value=0, interactive=False, precision=0)
58

59
        with gr.Row():
60
            prev_btn = gr.Button()
61
            next_btn = gr.Button()
62
            close_btn = gr.Button()
63

64
        with gr.Row():
65
            preview_samples = gr.JSON(interactive=False)
66

67
    dataset.change(
68
        can_preview, [dataset_dir, dataset], [data_preview_btn], queue=False
69
    ).then(
70
        lambda: 0, outputs=[page_index], queue=False
71
    )
72
    data_preview_btn.click(
73
        get_preview,
74
        [dataset_dir, dataset, page_index],
75
        [preview_count, preview_samples, preview_box],
76
        queue=False
77
    )
78
    prev_btn.click(
79
        prev_page, [page_index], [page_index], queue=False
80
    ).then(
81
        get_preview,
82
        [dataset_dir, dataset, page_index],
83
        [preview_count, preview_samples, preview_box],
84
        queue=False
85
    )
86
    next_btn.click(
87
        next_page, [page_index, preview_count], [page_index], queue=False
88
    ).then(
89
        get_preview,
90
        [dataset_dir, dataset, page_index],
91
        [preview_count, preview_samples, preview_box],
92
        queue=False
93
    )
94
    close_btn.click(lambda: gr.update(visible=False), outputs=[preview_box], queue=False)
95
    return dict(
96
        data_preview_btn=data_preview_btn,
97
        preview_count=preview_count,
98
        page_index=page_index,
99
        prev_btn=prev_btn,
100
        next_btn=next_btn,
101
        close_btn=close_btn,
102
        preview_samples=preview_samples
103
    )
104