es-aces 0.2.0__py3-none-any.whl

aces/__main__.py

@@ -0,0 +1,76 @@
+"""Main script for end-to-end task querying."""
+
+import os
+from pathlib import Path
+
+import hydra
+import polars as pl
+from loguru import logger
+from omegaconf import DictConfig
+
+from . import config, predicates, query
+
+
+def load_using_directory(cfg: config.TaskExtractorConfig, path: str | Path) -> pl.DataFrame:
+    # currently only supports ESGPT dataset when loading from directory
+    return predicates.generate_predicates_df(cfg, path, standard="ESGPT")
+
+
+def load_using_file(cfg: config.TaskExtractorConfig, path: str | Path) -> pl.DataFrame:
+    # .csv file
+    if path.suffix.lower() == ".csv":
+        return predicates.generate_predicates_df(cfg, path, standard="CSV")
+    # .parquet file (for MEDS)
+    elif path.suffix.lower() == ".parquet":
+        return predicates.generate_predicates_df(cfg, path, standard="MEDS")
+    else:
+        logger.error(f"Unsupported file type: {path.suffix}")
+        return
+
+
+@hydra.main(version_base=None, config_path="", config_name="")
+def main(cfg: DictConfig) -> None:
+    cfg = hydra.utils.instantiate(cfg, _convert_="all")
+
+    # Set output path
+    if "output_path" not in cfg or not cfg["output_path"]:
+        output_path = Path.cwd()
+    else:
+        output_path = Path(cfg["output_path"])
+
+    # load configuration
+    logger.info("Loading config...")
+    task_cfg_path = Path(cfg["config_path"])
+    task_cfg = config.TaskExtractorConfig.load(task_cfg_path)
+
+    # get predicates_df
+    logger.info("Loading data...")
+    data_path = Path(cfg["data_path"])
+    if not data_path.exists():
+        logger.error(f"{data_path} does not exist.")
+        return
+
+    if data_path.is_dir():
+        # load directory
+        logger.info("Directory provided, checking directory...")
+        predicates_df = load_using_directory(task_cfg, data_path)
+    else:
+        # load file
+        logger.info("File path provided, checking file...")
+        predicates_df = load_using_file(task_cfg, data_path)
+
+    # query results
+    result = query.query(task_cfg, predicates_df)
+
+    # save results
+    result = result.with_columns(
+        *[pl.col(col).struct.json_encode() for col in result.columns if "summary" in col]
+    ).write_csv(
+        os.path.join(output_path, "result.csv"),
+        include_header=True,
+        date_format="%m/%d/%Y %H:%M",
+    )
+
+
+if __name__ == "__main__":
+    main()