Source code for gunz_ml.management.scaffold

# -*- coding: utf-8 -*-
"""
Project Scaffolding and Boilerplate Generation.
"""

import os
import pathlib
import typing as t
import typer
from loguru import logger
from rich.console import Console

console = Console()

PYTHON_TEMPLATE = """# -*- coding: utf-8 -*-
\"\"\"
[PROJECT] {project_name} Reconstruction Script.
\"\"\"

__author__ = "{author}"
__email__ = "{email}"
__license__ = "Clear BSD"
__version__ = "1.0.0"

import os
import sys
import hydra
import lightning as L
from omegaconf import DictConfig, OmegaConf
from loguru import logger

def run_experiment(cfg: DictConfig):
    logger.info(f"Starting {project_name} with mode: {{cfg.args.mode}}")
    seed = cfg.args.get('seed', 42)
    L.seed_everything(seed)
    logger.success("Experiment complete.")

@hydra.main(version_base=None, config_path="configs", config_name="base_{project_name}")
def main(cfg: DictConfig) -> None:
    run_experiment(cfg)

if __name__ == "__main__":
    main()
"""

CONFIG_TEMPLATE = """# @package _global_
defaults:
  - /logger: mod_mlflow
  - /node: main
  - _self_

args:
  project_name: {project_name}
  mode: run
  seed: 42

optuna:
  n_trials: 100
  objectives:
    metric-spearmann_r: maximize
"""



def scaffold_project(
    project_name: str,
    interactive: bool = True,
    target_dir: t.Union[str, pathlib.Path] = "."
):
    target_path = pathlib.Path(target_dir)
    author = "Yeremia Gunawan Adhisantoso"
    email = "adhisant@tnt.uni-hannover.de"
    
    if interactive:
        author = typer.prompt("Project Author Name", default=author)
        email = typer.prompt("Project Author Email", default=email)

    py_content = PYTHON_TEMPLATE.format(
        project_name=project_name,
        author=author,
        email=email
    )
    cfg_content = CONFIG_TEMPLATE.format(project_name=project_name)

    py_file = target_path / f"{project_name}.py"
    if not py_file.exists() or typer.confirm(f"Overwrite {py_file}?", default=False):
        py_file.write_text(py_content)

    cfg_dir = target_path / "configs"
    cfg_dir.mkdir(exist_ok=True)
    base_cfg = cfg_dir / f"base_{project_name}.yaml"
    base_cfg.write_text(cfg_content)
    
    (cfg_dir / "exp" / project_name).mkdir(parents=True, exist_ok=True)
    
    console.print(f"[bold green]SUCCESS:[/bold green] Scaffolded '{project_name}'.")