mirror of
https://github.com/HannesStark/boltzgen.git
synced 2026-06-04 20:04:24 +08:00
54 lines
1.8 KiB
Python
Executable File
54 lines
1.8 KiB
Python
Executable File
import sys
|
|
from typing import List
|
|
import hydra
|
|
import omegaconf
|
|
|
|
|
|
from boltzgen.task.task import Task
|
|
|
|
|
|
def main(config: str, args: List) -> None:
|
|
"""
|
|
This is just a wrapper for running the .run() function of our `Task` class.
|
|
If you run the pipeline (for example via `boltzgen run design_spec.yaml ...`) then this function reads the yaml files of the individual pipeline steps and executes the pipeline steps.
|
|
|
|
The possible tasks are:
|
|
- Train (GPU: BoltzGen diffusion model or inverse folding model training)
|
|
- Predict (GPU: Running BoltzGen diffusion, inverse folding, refolding, designfolding, or affinity prediction)
|
|
- Analyze (CPU: Compute CPU Metrics and aggregate metrics from GPU steps)
|
|
- Filter (CPU: Very fast (20s) computes ranking and writes final output files)
|
|
|
|
The files for these are:
|
|
- src/boltzgen/task/train/train.py
|
|
- src/boltzgen/task/predict/predict.py
|
|
- src/boltzgen/task/analyze/analyze.py
|
|
- src/boltzgen/task/filter/filter.py
|
|
|
|
Parameters
|
|
----------
|
|
config : str
|
|
Path to the configuration yaml file. The yaml file contains something like `_target_: boltzgen.task.predict.predict.Predict` at the beginning which tells it which Task class to run
|
|
args : List
|
|
List of arguments to override the configuration.
|
|
"""
|
|
# Load the configuration
|
|
args = omegaconf.OmegaConf.from_dotlist(args)
|
|
config = omegaconf.OmegaConf.load(config)
|
|
config = omegaconf.OmegaConf.merge(config, args)
|
|
|
|
# Instantiate the task
|
|
task = hydra.utils.instantiate(config)
|
|
|
|
if not isinstance(task, Task):
|
|
msg = "Config must be an instance of Task."
|
|
raise TypeError(msg)
|
|
|
|
# Run the task
|
|
task.run(config)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
config = sys.argv[1]
|
|
args = sys.argv[2:]
|
|
main(config, args)
|