import yaml YAML_PATH = "./config.yaml" class Dumper(yaml.Dumper): def increase_indent(self, flow=False, *args, **kwargs): return super().increase_indent(flow=flow, indentless=False) # read scanners from yaml file # return a list of scanners def read_scanners(path): scanners = [] with open(path, "r") as f: config = yaml.load(f, Loader=yaml.FullLoader) scanners = config.get("detectors", None) return scanners # convert a list of scanners to yaml file def write_scanners(scanners): with open(YAML_PATH, "r") as f: config = yaml.load(f, Loader=yaml.FullLoader) config["detectors"] = scanners with open(YAML_PATH, "w") as f: # save scanners to detectors in yaml yaml.dump(config, f, Dumper=Dumper) # read model_type from yaml file def read_inference_type(path): inference_type = "" with open(path, "r") as f: config = yaml.load(f, Loader=yaml.FullLoader) inference_type = config.get("inference_type", None) return inference_type # write model_type to yaml file def write_inference_type(use_inference): with open(YAML_PATH, "r") as f: config = yaml.load(f, Loader=yaml.FullLoader) if use_inference: config["inference_type"] = 'hf_inference_api' else: config["inference_type"] = 'hf_pipeline' with open(YAML_PATH, "w") as f: # save inference_type to inference_type in yaml yaml.dump(config, f, Dumper=Dumper) # convert column mapping dataframe to json def convert_column_mapping_to_json(df, label=""): column_mapping = {} column_mapping[label] = [] for _, row in df.iterrows(): column_mapping[label].append(row.tolist()) return column_mapping