def main()

in florence2-VQA/src_train/train_mlflow.py [0:0]


def main(args):
    
    # Load model
    load_model(args.model_name_or_path)

    # Load datasets
    dataset = load_from_disk(args.train_dir)
    train_dataset = DocVQADataset(dataset['train'])
    val_dataset = DocVQADataset(dataset['validation'])

    # Train model
    train_model(args, train_dataset, val_dataset)