File size: 6,591 Bytes
e17c9f2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
r"""_summary_
-*- coding: utf-8 -*-

Module : configs.config

File Name : config.py

Description : Load the config file, which supports referencing other configuration files. If a circular reference occurs, an exception will be thrown
   
Creation Date : 2024-08-18

Author : Frank Kang([email protected])
"""
import pathlib
import json

import os
import warnings

from typing import Union, Any, IO
from omegaconf import OmegaConf, DictConfig, ListConfig

from .utils import get_dir

INCLUDE_KEY = 'include'


def get_api_aliases(llms_api, sum_api, gen_api):
    if sum_api is None:
        if llms_api is not None:
            sum_api = llms_api
        else:
            sum_api = 'ZhipuAI'

    if gen_api is None:
        if llms_api is not None:
            gen_api = llms_api
        else:
            gen_api = 'OpenAI'

    return sum_api, gen_api


def check_api_alias(config, api):
    api = api.lower()
    for k in config.keys():
        if k.lower() == api:
            return k
    return None


def update_config_with_api_aliases(config, llms_api, sum_api, gen_api):
    sum_api, gen_api = get_api_aliases(llms_api, sum_api, gen_api)
    sum_api_found = check_api_alias(config, sum_api)
    if sum_api_found is None:
        raise KeyError('{} cannot match any llms api in config'.format(sum_api))
    gen_api_found = check_api_alias(config, gen_api)
    if gen_api_found is None:
        raise KeyError('{} cannot match any llms api in config'.format(gen_api))
    config.used_llms_apis = {'summarization': sum_api_found, 'generation': gen_api_found}


class ConfigReader:
    """_summary_
    Load the config file, which supports referencing other configuration files. If a circular reference occurs, an exception will be thrown

    for examples:
    ```
    config = ConfigReader.load(file)
    ```
    """

    def __init__(
        self,
        file_: Union[str, pathlib.Path, IO[Any]],
        included: set | None = None
    ) -> None:
        """_summary_

        Args:
            file_ (Union[str, pathlib.Path, IO[Any]]): config
            included (set | None, optional): Include config file. Defaults to None.

        Raises:
            FileNotFoundError: If the configuration file cannot be found
            RecursionError: If there is a loop include
        """
        fname = ''
        self.included = included if included is not None else set()
        if isinstance(file_, str):
            fname = file_
            if not os.path.exists(fname):
                template_path = '{}.template'.format(fname)
                if os.path.exists(template_path):
                    with open(fname, 'w', encoding='utf8') as wf:
                        with open(template_path, 'r', encoding='utf8') as rf:
                            wf.write(rf.read())
                    warnings.warn(
                        'cannot find file {}. Auto generate from {}'.format(
                            fname, template_path))
                else:
                    raise FileNotFoundError(
                        'cannot find file {}'.format(fname))
        else:
            fname = file_.name

        suffix = fname.split('.')[-1]
        if suffix == 'yaml':
            config = OmegaConf.load(fname)
        elif suffix == 'json':
            if isinstance(file_, (str, IO[Any])):
                with open(file_, 'r', encoding='utf8') as f:
                    config = json.load(f)
            else:
                config = json.load(file_)
            config = DictConfig(config)
        if fname not in self.included:
            self.included.add(fname)
        else:
            raise RecursionError()
        self.__config = config
        self.complied = False

    def complie(self, config: DictConfig | None = None):
        """_summary_

        Resolve config to make include effective

        Args:
            config (DictConfig | None, optional): dict config. Defaults to None.

        Raises:
            RecursionError: If there is a loop include
        """
        modify_flag = False
        if config is None:
            config = self.__config
            modify_flag = True

        include_item = None

        if INCLUDE_KEY in config.keys():
            include_value = config.get(INCLUDE_KEY)
            if isinstance(include_value, (list, ListConfig)):
                include_item = [get_dir(p) for p in include_value]
            else:
                include_item = get_dir(include_value)
        for key in config.keys():
            value = config.get(key)
            if isinstance(value, DictConfig):
                self.complie(value)

        if include_item is not None:
            if isinstance(include_item, str):
                included = self.included.copy()
                if include_item in included:
                    print(include_item, included)
                    raise RecursionError()
                included.add(include_item)
                config.merge_with(ConfigReader.load(include_item, included))

            else:
                for item in include_item:
                    included = self.included.copy()
                    if item in included:
                        print(include_item, included)
                        raise RecursionError()
                    config.merge_with(ConfigReader.load(item, included))
                    included.add(item)

        if modify_flag:
            self.complied = True

    @property
    def config(self) -> DictConfig:
        """_summary_

        Obtain parsed dict config

        Returns:
            DictConfig: parsed dict config
        """
        if not self.complied:
            self.complie()
        return self.__config

    @staticmethod
    def load(
        file_: Union[str, pathlib.Path, IO[Any]],
        included: set | None = None,
        **kwargs
    ) -> DictConfig:
        """_summary_

        Class method loading configuration file

        Args:
            file_ (Union[str, pathlib.Path, IO[Any]]): config
            included (set | None, optional): Include config file. Defaults to None.

        Returns:
            DictConfig: parsed dict config
        """
        config = ConfigReader(file_, included).config
        if 'llms_api' in kwargs and 'sum_api' in kwargs and 'gen_api' in kwargs:
            update_config_with_api_aliases(config, kwargs['llms_api'], kwargs['sum_api'], kwargs['gen_api'])
            del kwargs['llms_api']
            del kwargs['sum_api']
            del kwargs['gen_api']
        for k, v in kwargs.items():
            config[k] = v
        return config