aurora

1
from typing import TYPE_CHECKING, Dict, List, Set
2

3
if TYPE_CHECKING:
4
    from gradio.components import Component
5

6

7
class Manager:
8

9
    def __init__(self) -> None:
10
        self.all_elems: Dict[str, Dict[str, "Component"]] = {}
11

12
    def get_elem_by_name(self, name: str) -> "Component":
13
        r"""
14
        Example: top.lang, train.dataset
15
        """
16
        tab_name, elem_name = name.split(".")
17
        return self.all_elems[tab_name][elem_name]
18

19
    def get_base_elems(self) -> Set["Component"]:
20
        return {
21
            self.all_elems["top"]["lang"],
22
            self.all_elems["top"]["model_name"],
23
            self.all_elems["top"]["model_path"],
24
            self.all_elems["top"]["checkpoints"],
25
            self.all_elems["top"]["finetuning_type"],
26
            self.all_elems["top"]["quantization_bit"],
27
            self.all_elems["top"]["template"],
28
            self.all_elems["top"]["flash_attn"],
29
            self.all_elems["top"]["shift_attn"],
30
            self.all_elems["top"]["rope_scaling"]
31
        }
32

33
    def list_elems(self) -> List["Component"]:
34
        return [elem for elems in self.all_elems.values() for elem in elems.values()]
35