| 7 | ZH_RAP_LORA_ROOT_DIR = "examples/zh_rap_lora/input_params" |
| 8 | |
| 9 | class DataSampler: |
| 10 | def __init__(self, root_dir=DEFAULT_ROOT_DIR): |
| 11 | self.root_dir = root_dir |
| 12 | self.input_params_files = list(Path(self.root_dir).glob("*.json")) |
| 13 | self.zh_rap_lora_input_params_files = list(Path(ZH_RAP_LORA_ROOT_DIR).glob("*.json")) |
| 14 | self.zh_rap_lora_input_params_files += list(Path(ZH_RAP_LORA_ROOT_DIR).glob("*.json")) |
| 15 | |
| 16 | def load_json(self, file_path): |
| 17 | with open(file_path, "r", encoding="utf-8") as f: |
| 18 | return json.load(f) |
| 19 | |
| 20 | def sample(self, lora_name_or_path=None): |
| 21 | if lora_name_or_path is None or lora_name_or_path == "none": |
| 22 | json_path = random.choice(self.input_params_files) |
| 23 | json_data = self.load_json(json_path) |
| 24 | else: |
| 25 | json_path = random.choice(self.zh_rap_lora_input_params_files) |
| 26 | json_data = self.load_json(json_path) |
| 27 | # Update the lora_name in the json_data |
| 28 | json_data["lora_name_or_path"] = lora_name_or_path |
| 29 | |
| 30 | return json_data |