Forge is a research library designed for representational learning in combinatorial optimization. It provides tools for generating embeddings from MIP instances, pre-training models on these embeddings, and fine-tuning them for specific tasks such as predicting integral gap, search guidance, backdoor prediction, and solver configuration.
# Install the library
pip install forge-mip
# Generate MIP Embeddings from the Hugging Face pre-trained Forge model and save the output
# Export your Hugging Face Token, if not already set in your environment
# export HF_TOKEN=<your_hugging_face_token>
forge --input_mips ./data/instances/ --input_mip_instances_file ./data/configs/test.txt --output_mip_to_embeddings_pkl ./models/mip_to_embeddings.pkl
# Generate MIP Embeddings from a local pre-trained Forge model and save the output
forge --train_config_yaml ./forge/configs/train_config.yaml --input_forge_pkl ./models/forge_pretrained.pkl --input_mips ./data/instances/ --input_mip_instances_file ./data/configs/test.txt --output_mip_to_embeddings_pkl ./models/mip_to_embeddings.pklNote: When no local model is found, forge downloads the pre-trained Forge model from Hugging Face and saves it to the package's models directory — for a pip install-ed package that is <site-packages>/models/forge_pretrained.pkl, not ./models/. To use a model stored in the repo, place it at ./models/forge_pretrained.pkl and pass --input_forge_pkl ./models/forge_pretrained.pkl.
Access Request: The pretrained Forge model is gated, please request access first. Without an HF token with approved access, the load from Hugging Face will fail.
| Functionality | Description |
|---|---|
| Generate MIP Info | Build and serialize MIPInfo objects from raw MIP instances for reuse in downstream pre-training. |
| Pre-Train Forge | Pre-train Forge on MIP instances and their MIPInfo and save a pretrained model checkpoint. |
| Generate Embeddings | Generate per-instance embeddings from a pretrained Forge model (MIPEmbeddings). |
| Fine-Tune Integral Gap | Fine-tune Forge for integral-gap prediction on labeled GapInfo data. |
| Predict Integral Gap | Run inference with a fine-tuned model to predict LP/MIP gap information (GapInfo). |
| Fine-Tune Variable Probabilities | Fine-tune Forge for variable 0/1 probability prediction on labeled TripletInfo data. |
| Predict Variable Probabilities | Run inference with a fine-tuned model to predict variable 0/1 probabilities (HintInfo). |
from forge.embeddings import Forge
from forge.pipeline import mip_to_mipinfo
from forge.utils import Constants
# Forge model with its pre-trained configuration
forge = Forge(train_config_yaml="./forge/configs/train_config.yaml")
# Generate MIP info object for a set of given mip instances
# The output mip_to_mipinfo pickle is stored as output_mip_to_mipinfo_pkl
# The mip_to_mipinfo pkl can be re-used in pretrain() with input_mip_to_mipinfo_pkl flag
# - mip_to_mipinfo maps mip instance to a mipinfo object, Dict[str, MIPInfo], containing:
# - instance_name: str, the name of the MIP instance
# - feature_tensor: torch.Tensor, the feature tensor for the MIP instance (num_cons + num_vars, feat_dim=10)
# - num_cons: int, the number of constraints in the MIP instance
# - num_vars: int, the number of variables in the MIP instance
# - edge_index: torch.Tensor, (2, E) edges from source (constraint) to target (variable) nodes
# - edge_weight: torch.Tensor, (E,), weights of the edges
# Pretraining log is stored in output_log_file with loss curves and training details
mip_to_mipinfo(forge=forge,
input_mip_folder="./data/instances/",
input_mip_instances_file="./data/configs/test.txt",
output_mip_to_mipinfo_pkl="./models/test_mip_to_mipinfo.pkl",
relaxation_list=[0.05, 0.01],
num_parallel_workers=1)cd forge
python -m scripts.mip_to_mipinfo --train_config_yaml ./forge/configs/train_config.yaml --input_mip_folder ./data/instances/ --input_mip_instances_file ./data/configs/all.txt --output_mip_to_mipinfo_pkl ./models/mip_to_mipinfo.pkl --relaxation_list 0.05 0.01 --num_parallel_workers 1from forge.embeddings import Forge
from forge.pipeline import pretrain
# Forge model with its pre-training configuration
forge = Forge(train_config_yaml="./forge/configs/train_config.yaml")
# Pretrain Forge on a set of MIP instances in the given input folder
# The pretrained model pickle is stored as output_forge_pretrained_pkl
# The intermediate mip_to_mipinfo pickle is stored as output_mip_to_mipinfo_pkl
# The mip_to_mipinfo pkl can be reused with input_mip_to_mipinfo_pkl flag to skip MIP parsing in future pre-training
# - mip_to_mipinfo maps mip instance to a mipinfo object, Dict[str, MIPInfo], containing:
# - instance_name: str, the name of the MIP instance
# - feature_tensor: torch.Tensor, the feature tensor for the MIP instance (num_cons + num_vars, feat_dim=10)
# - num_cons: int, the number of constraints in the MIP instance
# - num_vars: int, the number of variables in the MIP instance
# - edge_index: torch.Tensor, (2, E) edges from source (constraint) to target (variable) nodes
# - edge_weight: torch.Tensor, (E,), weights of the edges
# Pretraining log is stored in output_log_file with loss curves and training details
pretrain(forge=forge,
input_mip_folder="./data/instances/",
input_mip_instances_file="data/configs/all.txt",
output_mip_to_mipinfo_pkl="./models/pretrain_clusters_mip_to_mipinfo.pkl",
output_forge_pretrained_pkl="./models/forge_pretrained.pkl",
output_log_file="./models/forge_pretrained.log")cd forge
python -m scripts.pretrain --train_config_yaml ./forge/configs/train_config.yaml --input_mip_folder ./data/instances/ --input_mip_instances_file ./data/configs/all.txt --relaxation_list 0.05 0.01 --output_mip_to_mipinfo_pkl ./models/pretrain_clusters_mip_to_mipinfo.pkl --output_forge_pretrained_pkl ./models/forge_pretrained.pkl --output_log_file ./models/forge_pretrained.logfrom forge.embeddings import Forge
from forge.pipeline import mip_to_embeddings
from forge.utils import Constants
# Forge model with its pre-trained configuration
forge = Forge(train_config_yaml="./forge/configs/train_config.yaml")
# Generate embeddings dictionary for MIPs in the input folder
# Use the trained Forge model stored in input_forge_pkl of type model_type
# The output mip_to_embeddings pickle is stored as output_mip_to_embeddings_pkl
# Each MIP instance is mapped to a MIPEmbeddings object, Dict[str, MIPEmbeddings], containing:
# - instance_embedding: np.ndarray (codebook_size)
# - embedding_of_constraint: torch.Tensor (num_constraints, codebook_dim)
# - embedding_of_variable: torch.Tensor (num_variables, codebook_dim)
mip_to_embeddings_dict = mip_to_embeddings(forge=forge,
input_mips="./data/instances/",
input_mip_instances_file="./data/configs/test.txt",
input_forge_pkl="./models/forge_pretrained.pkl",
model_type=Constants.FORGE_PRE_TRAIN,
instance_embedding_only=False,
output_mip_to_embeddings_pkl="./models/mip_to_embeddings.pkl")cd forge
python -m scripts.mip_to_embeddings --train_config_yaml ./forge/configs/train_config.yaml --input_forge_pkl ./models/forge_pretrained.pkl --input_mips ./data/instances/ --input_mip_instances_file ./data/configs/test.txt --output_mip_to_embeddings_pkl ./models/mip_to_embeddings.pklfrom forge.embeddings import Forge
from forge.pipeline import finetune_integral_gap
from forge.utils import Constants
# Forge model with its pre-trained configuration
forge = Forge(train_config_yaml="./forge/configs/train_config.yaml")
# Fine-tune Forge to predict integral gaps
finetune_integral_gap(forge=forge,
input_forge_pkl="./models/forge_pretrained.pkl",
model_type=Constants.FORGE_FINE_TUNE_INTEGRAL_GAP,
input_mip_folder="./data/instances/",
input_mip_instances_file="data/configs/fine_tune_integral_gap.txt",
output_forge_finetuned_pkl="./models/forge_integral_gap.pkl",
output_mip_to_gapinfo_pkl="./models/mip_to_gapinfo.pkl",
num_parallel_workers=5)cd forge
python -m scripts.finetune_integral_gap --train_config_yaml ./forge/configs/train_config.yaml --input_forge_pkl ./models/forge_pretrained.pkl --input_mip_folder ./data/instances/ --input_mip_instances_file ./data/configs/fine_tune_integral_gap.txt --output_forge_finetuned_pkl ./models/forge_integral_gap.pkl --output_mip_to_gapinfo_pkl ./models/mip_to_gapinfo.pklfrom forge.embeddings import Forge
from forge.pipeline import mip_to_gapinfo
from forge.utils import Constants
# Forge model with its pre-trained configuration
forge = Forge(train_config_yaml="./forge/configs/train_config.yaml")
# Predict integral gaps
# Each MIP instance is mapped to a GapInfo object, Dict[str, GapInfo], containing:
# - lp_obj: the true objective value of the lp relaxation solution
# - lp_sol: the true lp relaxation solution
# - mip_obj: the predicted objective value of the mip solution
# - mip_sol: None, there is no solution, only gap prediction
# - gap_ratio: float, the predicted ratio between lp and mip
mip_to_gapinfo_dict = mip_to_gapinfo(forge=forge,
input_forge_pkl="./models/forge_integral_gap.pkl",
model_type=Constants.FORGE_FINE_TUNE_INTEGRAL_GAP,
input_mips="./data/instances/",
input_mip_instances_file="./data/configs/test_integral_gap.txt",
output_mip_to_gapinfo_pkl="./models/mip_to_gapinfo.pkl",
problem_type="CA")cd forge
python -m scripts.mip_to_gapinfo --train_config_yaml ./forge/configs/train_config.yaml --input_forge_pkl ./models/forge_integral_gap.pkl --input_mips ./data/instances/ --input_mip_instances_file ./data/configs/test_integral_gap.txt --output_mip_to_gapinfo_pkl ./models/mip_to_gapinfo.pkl --problem_type CAfrom forge.embeddings import Forge
from forge.pipeline import finetune_variable_proba
from forge.utils import Constants
# Forge model with its pre-trained configuration
forge = Forge(train_config_yaml="./forge/configs/train_config.yaml")
# Fine-tune Forge to predict variable probabilities
finetune_variable_proba(forge=forge,
input_forge_pkl="./models/forge_pretrained.pkl",
model_type=Constants.FORGE_FINE_TUNE_VARIABLE_PROBA,
input_mip_folder="./data/instances/",
input_mip_instances_file="data/configs/fine_tune_variable_proba.txt",
output_forge_finetuned_pkl="./models/forge_variable_proba.pkl",
output_mip_to_tripletinfo_pkl="./models/output_mip_to_tripletinfo.pkl",
triplet_time_limit=300,
triplet_num_solutions=5)cd forge
python -m scripts.finetune_variable_proba --train_config_yaml ./forge/configs/train_config.yaml --input_forge_pkl ./models/forge_pretrained.pkl --input_mip_folder ./data/instances/ --input_mip_instances_file ./data/configs/fine_tune_variable_proba.txt --output_forge_finetuned_pkl ./models/forge_variable_proba.pkl --output_mip_to_tripletinfo_pkl ./models/output_mip_to_tripletinfo.pklfrom forge.embeddings import Forge
from forge.pipeline import mip_to_hint
from forge.utils import Constants
# Forge model with its pre-trained configuration
forge = Forge(train_config_yaml="./forge/configs/train_config.yaml")
# Predict variable 0/1 probabilities
# Each MIP instance is mapped to a HintInfo object, Dict[str, HintInfo], containing:
# - hint_ones: variable indices to hint as 1
# - hint_zeros: variable indices to hint as 0
# - hint_pri_ones: priority ranks for the 1-hints (higher = more confident)
# - hint_pri_zeros: priority ranks for the 0-hints (higher = more confident)
mip_to_hintinfo_dict = mip_to_hint(forge=forge,
input_forge_pkl="./models/forge_variable_proba.pkl",
model_type=Constants.FORGE_FINE_TUNE_VARIABLE_PROBA,
input_mips="./data/instances/",
input_mip_instances_file="./data/configs/test_variable_proba.txt",
output_mip_to_hintinfo_pkl="./models/mip_to_hintinfo.pkl",
problem_type="CA")cd forge
python -m scripts.mip_to_hint --train_config_yaml ./forge/configs/train_config.yaml --input_forge_pkl ./models/forge_variable_proba.pkl --input_mips ./data/instances/ --input_mip_instances_file ./data/configs/test_variable_proba.txt --output_mip_to_hintinfo_pkl ./models/mip_to_hintinfo.pkl --problem_type CAForge requires Python 3.10 or newer and can be installed via pip install forge-mip.
git clone https://github.com/skadio/forge.git
cd forge
pip install build # if build is not installed
python -m build
pip install dist/forge_mip-X.X.X-py3-none-any.whl
$ git clone https://github.com/skadio/forge.git
$ cd forge
$ python -m unittest discover tests
Please submit bug reports and feature requests as Issues.
We would like to thank Modal for their generous support through the provision of academic credits and computational infrastructure, which were instrumental in training the Forge model used in this research.
Forge is licensed under the Apache License 2.0.