gridfm-graphkit
This library is brought to you by the GridFM team to train, finetune and interact with a foundation model for the electric power grid.
You can install gridfm-graphkit
directly from PyPI:
pip install gridfm-graphkit
To contribute or develop locally, clone the repository and install in editable mode:
git clone git@github.com:gridfm/gridfm-graphkit.git
cd gridfm-graphkit
python -m venv venv
source venv/bin/activate
pip install -e .
For documentation generation and unit testing, install with the optional dev
and test
extras:
pip install -e .[dev,test]
An interface to train, fine-tune, and evaluate GridFM models using configurable YAML files and MLflow tracking.
gridfm_graphkit <command> [OPTIONS]
Available commands:
train
– Train a new model from scrathcfinetune
– Fine-tune an existing pre-trained modelevaluate
– Evaluate model performance on a datasetpredict
– Run inference and save predictions
gridfm_graphkit train --config path/to/config.yaml
Argument | Type | Description | Default |
---|---|---|---|
--config |
str |
Required. Path to the training configuration YAML file. | None |
--exp_name |
str |
Optional. MLflow experiment name. | timestamp |
--run_name |
str |
Optional. MLflow run name. | run |
--log_dir |
str |
Optional. MLflow logging directory. | mlruns |
--data_path |
str |
Optional. Root dataset directory. | data |
Standard Training:
gridfm_graphkit train --config examples/config/case30_ieee_base.yaml --data_path examples/data
gridfm_graphkit finetune --config path/to/config.yaml --model_path path/to/model.pth
Argument | Type | Description | Default |
---|---|---|---|
--config |
str |
Required. Fine-tuning configuration file. | None |
--model_path |
str |
Required. Path to a pre-trained model file. | None |
--exp_name |
str |
MLflow experiment name. | timestamp |
--run_name |
str |
MLflow run name. | run |
--log_dir |
str |
MLflow logging directory. | mlruns |
--data_path |
str |
Root dataset directory. | data |
gridfm_graphkit evaluate --config path/to/eval.yaml --model_path path/to/model.pth
Argument | Type | Description | Default |
---|---|---|---|
--config |
str |
Required. Path to evaluation config. | None |
--model_path |
str |
Path to the trained model file. | None |
--exp_name |
str |
MLflow experiment name. | timestamp |
--run_name |
str |
MLflow run name. | run |
--log_dir |
str |
MLflow logging directory. | mlruns |
--data_path |
str |
Dataset directory. | data |
gridfm_graphkit predict --config path/to/config.yaml --model_path path/to/model.pth
Argument | Type | Description | Default |
---|---|---|---|
--config |
str |
Required. Path to prediction config file. | None |
--model_path |
str |
Path to the trained model file. | None |
--exp_name |
str |
MLflow experiment name. | timestamp |
--run_name |
str |
MLflow run name. | run |
--log_dir |
str |
MLflow logging directory. | mlruns |
--data_path |
str |
Dataset directory. | data |
--output_path |
str |
Directory where predictions are saved. | data |