optimizer_config.schema.json

{
    "$defs": {
        "CrossEncoderConfig": {
            "additionalProperties": false,
            "properties": {
                "model_name": {
                    "default": "cross-encoder/ms-marco-MiniLM-L6-v2",
                    "description": "Name of the hugging face model.",
                    "title": "Model Name",
                    "type": "string"
                },
                "batch_size": {
                    "default": 32,
                    "description": "Batch size for model inference.",
                    "exclusiveMinimum": 0,
                    "title": "Batch Size",
                    "type": "integer"
                },
                "device": {
                    "anyOf": [
                        {
                            "type": "string"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "default": null,
                    "description": "Torch notation for CPU or CUDA.",
                    "title": "Device"
                },
                "bf16": {
                    "default": false,
                    "description": "Whether to use mixed precision training (not all devices support this).",
                    "title": "Bf16",
                    "type": "boolean"
                },
                "fp16": {
                    "default": false,
                    "description": "Whether to use mixed precision training (not all devices support this).",
                    "title": "Fp16",
                    "type": "boolean"
                },
                "tokenizer_config": {
                    "$ref": "#/$defs/TokenizerConfig"
                },
                "trust_remote_code": {
                    "default": false,
                    "description": "Whether to trust the remote code when loading the model.",
                    "title": "Trust Remote Code",
                    "type": "boolean"
                },
                "revision": {
                    "anyOf": [
                        {
                            "type": "string"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "default": null,
                    "description": "Revision from HF repo",
                    "title": "Revision"
                },
                "train_head": {
                    "default": false,
                    "description": "Whether to train the head of the model. If False, LogReg will be trained.",
                    "title": "Train Head",
                    "type": "boolean"
                }
            },
            "title": "CrossEncoderConfig",
            "type": "object"
        },
        "DataConfig": {
            "additionalProperties": false,
            "description": "Configuration for the data used in the optimization process.",
            "properties": {
                "scheme": {
                    "default": "ho",
                    "description": "Validation scheme to use.",
                    "enum": [
                        "ho",
                        "cv"
                    ],
                    "title": "Scheme",
                    "type": "string"
                },
                "n_folds": {
                    "default": 3,
                    "description": "Number of folds in cross-validation.",
                    "exclusiveMinimum": 0,
                    "title": "N Folds",
                    "type": "integer"
                },
                "validation_size": {
                    "default": 0.2,
                    "description": "Fraction of train samples to allocate for validation (if input dataset doesn't contain validation split). If `is_few_shot_train` is True, this value will be ignored.",
                    "maximum": 1,
                    "minimum": 0,
                    "title": "Validation Size",
                    "type": "number"
                },
                "separation_ratio": {
                    "anyOf": [
                        {
                            "maximum": 1,
                            "minimum": 0,
                            "type": "number"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "default": null,
                    "description": "Set to float to prevent data leak between scoring and decision nodes.",
                    "title": "Separation Ratio"
                },
                "is_few_shot_train": {
                    "default": false,
                    "description": "Whether to use few-shot training.",
                    "title": "Is Few Shot Train",
                    "type": "boolean"
                },
                "examples_per_intent": {
                    "default": 8,
                    "description": "Number of examples per intent for few-shot validation. If None, all examples will be used.",
                    "exclusiveMinimum": 0,
                    "title": "Examples Per Intent",
                    "type": "integer"
                }
            },
            "title": "DataConfig",
            "type": "object"
        },
        "EmbedderConfig": {
            "additionalProperties": false,
            "description": "Base class for embedder configurations.",
            "properties": {
                "default_prompt": {
                    "anyOf": [
                        {
                            "type": "string"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "default": null,
                    "description": "Default prompt for the model. This is used when no task specific prompt is not provided.",
                    "title": "Default Prompt"
                },
                "classification_prompt": {
                    "anyOf": [
                        {
                            "type": "string"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "default": null,
                    "description": "Prompt for classifier.",
                    "title": "Classification Prompt"
                },
                "cluster_prompt": {
                    "anyOf": [
                        {
                            "type": "string"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "default": null,
                    "description": "Prompt for clustering.",
                    "title": "Cluster Prompt"
                },
                "sts_prompt": {
                    "anyOf": [
                        {
                            "type": "string"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "default": null,
                    "description": "Prompt for finding most similar sentences.",
                    "title": "Sts Prompt"
                },
                "query_prompt": {
                    "anyOf": [
                        {
                            "type": "string"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "default": null,
                    "description": "Prompt for query.",
                    "title": "Query Prompt"
                },
                "passage_prompt": {
                    "anyOf": [
                        {
                            "type": "string"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "default": null,
                    "description": "Prompt for passage.",
                    "title": "Passage Prompt"
                },
                "use_cache": {
                    "default": true,
                    "description": "Whether to use embeddings caching.",
                    "title": "Use Cache",
                    "type": "boolean"
                }
            },
            "title": "EmbedderConfig",
            "type": "object"
        },
        "HFModelConfig": {
            "additionalProperties": false,
            "properties": {
                "model_name": {
                    "default": "prajjwal1/bert-tiny",
                    "description": "Name of the hugging face repository with transformer model.",
                    "title": "Model Name",
                    "type": "string"
                },
                "batch_size": {
                    "default": 32,
                    "description": "Batch size for model inference.",
                    "exclusiveMinimum": 0,
                    "title": "Batch Size",
                    "type": "integer"
                },
                "device": {
                    "anyOf": [
                        {
                            "type": "string"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "default": null,
                    "description": "Torch notation for CPU or CUDA.",
                    "title": "Device"
                },
                "bf16": {
                    "default": false,
                    "description": "Whether to use mixed precision training (not all devices support this).",
                    "title": "Bf16",
                    "type": "boolean"
                },
                "fp16": {
                    "default": false,
                    "description": "Whether to use mixed precision training (not all devices support this).",
                    "title": "Fp16",
                    "type": "boolean"
                },
                "tokenizer_config": {
                    "$ref": "#/$defs/TokenizerConfig"
                },
                "trust_remote_code": {
                    "default": false,
                    "description": "Whether to trust the remote code when loading the model.",
                    "title": "Trust Remote Code",
                    "type": "boolean"
                },
                "revision": {
                    "anyOf": [
                        {
                            "type": "string"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "default": null,
                    "description": "Revision from HF repo",
                    "title": "Revision"
                }
            },
            "title": "HFModelConfig",
            "type": "object"
        },
        "HPOConfig": {
            "additionalProperties": false,
            "description": "Configuration for hyperparameter optimization using Optuna.\n\nFor more detailed information about the TPE sampler and its parameters,\nrefer to Optuna's documentation of `TPESampler\n<https://optuna.readthedocs.io/en/stable/reference/samplers/generated/optuna.samplers.TPESampler.html>`_,\n`study.optimize\n<https://optuna.readthedocs.io/en/stable/reference/generated/optuna.study.Study.html#optuna.study.Study.optimize>`_,\n`RandomSampler <https://optuna.readthedocs.io/en/stable/reference/samplers/generated/optuna.samplers.RandomSampler.html>`_.",
            "properties": {
                "sampler": {
                    "default": "tpe",
                    "description": "Optuna sampler. One of 'tpe', 'random'.",
                    "enum": [
                        "tpe",
                        "random"
                    ],
                    "title": "Sampler",
                    "type": "string"
                },
                "n_trials": {
                    "default": 15,
                    "description": "Number of trials to run in the optimization process. This is the total number of different hyperparameter combinations that will be evaluated.",
                    "title": "N Trials",
                    "type": "integer"
                },
                "timeout": {
                    "anyOf": [
                        {
                            "type": "number"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "default": null,
                    "description": "Time limit in seconds for the optimization process. If None, the optimization will run until n_trials is reached.",
                    "title": "Timeout"
                },
                "n_jobs": {
                    "default": 1,
                    "description": "Number of parallel jobs to run. Set to -1 to use all available CPU cores.",
                    "title": "N Jobs",
                    "type": "integer"
                },
                "n_startup_trials": {
                    "default": 10,
                    "description": "Number of initial trials to run using random sampling before switching to TPE algorithm. This helps in better initialization of the TPE algorithm.",
                    "title": "N Startup Trials",
                    "type": "integer"
                },
                "consider_prior": {
                    "default": true,
                    "description": "Whether to use Gaussian (normal) distribution as prior for integer and float parameter spaces. This helps in better initialization of the TPE algorithm's parameter distributions.",
                    "title": "Consider Prior",
                    "type": "boolean"
                },
                "prior_weight": {
                    "default": 1,
                    "description": "Weight of the prior distribution in the TPE algorithm. Higher values make the algorithm more conservative in exploring new regions.",
                    "title": "Prior Weight",
                    "type": "integer"
                },
                "n_ei_candidates": {
                    "default": 24,
                    "description": "Number of candidates to sample for expected improvement calculation in TPE algorithm. Higher values may lead to better exploration but slower optimization.",
                    "title": "N Ei Candidates",
                    "type": "integer"
                },
                "constant_liar": {
                    "default": false,
                    "description": "Whether to use constant liar strategy for parallel optimization. If True, the algorithm will penalize running trials to avoid suggesting parameter configurations that are too close to currently running trials.",
                    "title": "Constant Liar",
                    "type": "boolean"
                }
            },
            "title": "HPOConfig",
            "type": "object"
        },
        "LoggingConfig": {
            "additionalProperties": false,
            "description": "Configuration for the logging.",
            "properties": {
                "project_dir": {
                    "anyOf": [
                        {
                            "format": "path",
                            "type": "string"
                        },
                        {
                            "type": "string"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "default": null,
                    "description": "Path to the directory with different runs.",
                    "title": "Project Dir"
                },
                "run_name": {
                    "anyOf": [
                        {
                            "type": "string"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "default": null,
                    "description": "Name of the run. If None, a random name will be generated.",
                    "title": "Run Name"
                },
                "dump_modules": {
                    "default": false,
                    "description": "Whether to dump the modules or not",
                    "title": "Dump Modules",
                    "type": "boolean"
                },
                "clear_ram": {
                    "default": false,
                    "description": "Whether to clear the RAM after dumping the modules",
                    "title": "Clear Ram",
                    "type": "boolean"
                },
                "report_to": {
                    "anyOf": [
                        {
                            "items": {
                                "enum": [
                                    "wandb",
                                    "tensorboard",
                                    "codecarbon",
                                    "none"
                                ],
                                "type": "string"
                            },
                            "type": "array"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "default": [
                        "none"
                    ],
                    "description": "List of callbacks to report to. If None, no callbacks will be used",
                    "title": "Report To"
                },
                "log_interval_time": {
                    "default": 0.1,
                    "description": "Sampling interval for the system monitor in seconds for Wandb logger.",
                    "title": "Log Interval Time",
                    "type": "number"
                }
            },
            "title": "LoggingConfig",
            "type": "object"
        },
        "TokenizerConfig": {
            "additionalProperties": false,
            "properties": {
                "padding": {
                    "anyOf": [
                        {
                            "type": "boolean"
                        },
                        {
                            "enum": [
                                "longest",
                                "max_length",
                                "do_not_pad"
                            ],
                            "type": "string"
                        }
                    ],
                    "default": true,
                    "title": "Padding"
                },
                "truncation": {
                    "default": true,
                    "title": "Truncation",
                    "type": "boolean"
                },
                "max_length": {
                    "anyOf": [
                        {
                            "exclusiveMinimum": 0,
                            "type": "integer"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "default": null,
                    "description": "Maximum length of input sequences.",
                    "title": "Max Length"
                }
            },
            "title": "TokenizerConfig",
            "type": "object"
        }
    },
    "description": "Configuration for the optimization process.\n\nOne can use it to customize optimization beyond choosing different preset.\nInstantiate it and pass to :py:meth:`autointent.Pipeline.from_optimization_config`.",
    "properties": {
        "data_config": {
            "$ref": "#/$defs/DataConfig",
            "default": {
                "scheme": "ho",
                "n_folds": 3,
                "validation_size": 0.2,
                "separation_ratio": null,
                "is_few_shot_train": false,
                "examples_per_intent": 8
            }
        },
        "search_space": {
            "items": {
                "additionalProperties": true,
                "type": "object"
            },
            "title": "Search Space",
            "type": "array"
        },
        "logging_config": {
            "$ref": "#/$defs/LoggingConfig",
            "default": {
                "project_dir": null,
                "run_name": null,
                "dump_modules": false,
                "clear_ram": false,
                "report_to": [
                    "none"
                ],
                "log_interval_time": 0.1
            }
        },
        "embedder_config": {
            "$ref": "#/$defs/EmbedderConfig"
        },
        "cross_encoder_config": {
            "$ref": "#/$defs/CrossEncoderConfig",
            "default": {
                "model_name": "cross-encoder/ms-marco-MiniLM-L6-v2",
                "batch_size": 32,
                "device": null,
                "bf16": false,
                "fp16": false,
                "tokenizer_config": {
                    "max_length": 512,
                    "padding": true,
                    "truncation": true
                },
                "trust_remote_code": false,
                "revision": null,
                "train_head": false
            }
        },
        "transformer_config": {
            "$ref": "#/$defs/HFModelConfig",
            "default": {
                "model_name": "prajjwal1/bert-tiny",
                "batch_size": 32,
                "device": null,
                "bf16": false,
                "fp16": false,
                "tokenizer_config": {
                    "max_length": null,
                    "padding": true,
                    "truncation": true
                },
                "trust_remote_code": false,
                "revision": "refs/pr/16"
            }
        },
        "hpo_config": {
            "$ref": "#/$defs/HPOConfig",
            "default": {
                "sampler": "tpe",
                "n_trials": 15,
                "timeout": null,
                "n_jobs": 1,
                "n_startup_trials": 10,
                "consider_prior": true,
                "prior_weight": 1,
                "n_ei_candidates": 24,
                "constant_liar": false
            }
        },
        "seed": {
            "default": 42,
            "exclusiveMinimum": 0,
            "title": "Seed",
            "type": "integer"
        }
    },
    "required": [
        "search_space"
    ],
    "title": "OptimizationConfig",
    "type": "object"
}