Spaces:
Running
Running
File size: 844 Bytes
9ff79dc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 |
from pathlib import Path
from colpali_engine.utils.train_colpali_engine_models import ColModelTrainingConfig, ColModelTraining
from colpali_engine.utils.gpu_stats import print_gpu_utilization
import typer
import configue
def main(config_file: Path) -> None:
print_gpu_utilization()
print("Loading config")
config = configue.load(config_file, sub_path="config")
print("Creating Setup")
if isinstance(config, ColModelTrainingConfig):
app = ColModelTraining(config)
else:
raise ValueError("Config must be of type ColModelTrainingConfig")
if config.run_train:
print("Training model")
app.train()
app.save(config_file=config_file)
if config.run_eval:
print("Running evaluation")
app.eval()
print("Done!")
if __name__ == "__main__":
typer.run(main)
|