257 lines
8.1 KiB
Python
257 lines
8.1 KiB
Python
# utils/config_loader.py
|
||
|
||
import argparse
|
||
import copy
|
||
import logging
|
||
from pathlib import Path
|
||
from typing import Dict, List, Sequence, Any
|
||
|
||
import yaml
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
_ALWAYS: Sequence[str] = (
|
||
# minimal required keys for the pipeline to run at all
|
||
"experiment_base_dir",
|
||
"human_profile_path",
|
||
"num_iterations",
|
||
"min_word_len_for_analysis",
|
||
"log_level",
|
||
"model_id",
|
||
)
|
||
|
||
_VLLM: Sequence[str] = (
|
||
"vllm_model_id",
|
||
"vllm_port",
|
||
"vllm_hf_token",
|
||
"vllm_cuda_visible_devices",
|
||
"vllm_gpu_memory_utilization",
|
||
"vllm_max_model_len",
|
||
"vllm_dtype",
|
||
"vllm_extra_args",
|
||
"vllm_env",
|
||
)
|
||
|
||
_GENERATION: Sequence[str] = (
|
||
"generation_api_key",
|
||
"generation_api_base_url", # needed if you do local or remote calls
|
||
"generation_model_id",
|
||
"generation_max_new_tokens",
|
||
"generation_threads",
|
||
"generation_max_prompts",
|
||
"generation_hf_dataset_name",
|
||
"generation_hf_dataset_split",
|
||
"generation_logging_level",
|
||
"generation_chat_template_model_id",
|
||
"generation_param_chunk_size",
|
||
"generation_param_top_logprobs_count",
|
||
"generation_param_temperature",
|
||
"generation_param_top_p",
|
||
"generation_param_top_k",
|
||
"generation_param_min_p",
|
||
"generation_param_timeout",
|
||
"generation_param_stop_sequences",
|
||
"generation_ngram_remove_stopwords",
|
||
"generation_ngram_language",
|
||
"generation_force_backtrack",
|
||
"generation_prompt_template",
|
||
"generation_system_prompt"
|
||
|
||
)
|
||
|
||
_NGRAM: Sequence[str] = (
|
||
"top_k_bigrams",
|
||
"top_k_trigrams",
|
||
"dict_bigrams_initial",
|
||
"dict_bigrams_subsequent",
|
||
"nodict_bigrams_initial",
|
||
"nodict_bigrams_subsequent",
|
||
"dict_trigrams_initial",
|
||
"dict_trigrams_subsequent",
|
||
"nodict_trigrams_initial",
|
||
"nodict_trigrams_subsequent",
|
||
"extra_ngrams_to_ban",
|
||
)
|
||
|
||
_SLOP: Sequence[str] = (
|
||
"min_phrase_freq_to_keep",
|
||
"top_n_initial_slop_ban",
|
||
"top_n_subsequent_slop_ban",
|
||
"extra_slop_phrases_to_ban",
|
||
"banned_slop_phrases_filename",
|
||
)
|
||
|
||
_OVERREP: Sequence[str] = (
|
||
"top_k_words_for_overrep_analysis",
|
||
"dict_overrep_initial",
|
||
"dict_overrep_subsequent",
|
||
"nodict_overrep_initial",
|
||
"nodict_overrep_subsequent",
|
||
)
|
||
|
||
_FINETUNE: Sequence[str] = (
|
||
"finetune_mode",
|
||
"finetune_ftpo_dataset",
|
||
"finetune_base_model_id",
|
||
"finetune_max_seq_length",
|
||
"finetune_load_in_4bit",
|
||
"finetune_lora_r",
|
||
"finetune_lora_alpha",
|
||
"finetune_lora_dropout",
|
||
"finetune_weight_decay",
|
||
"finetune_target_modules",
|
||
"finetune_gradient_checkpointing",
|
||
"finetune_chat_template",
|
||
"finetune_batch_size",
|
||
"finetune_gradient_accumulation_steps",
|
||
"finetune_warmup_ratio",
|
||
"finetune_num_epochs",
|
||
"finetune_learning_rate",
|
||
"finetune_auto_learning_rate",
|
||
"finetune_beta",
|
||
"finetune_output_dir_suffix",
|
||
"finetune_save_merged_16bit",
|
||
"finetune_save_gguf_q8_0",
|
||
"finetune_max_train_examples",
|
||
"finetune_cuda_visible_devices",
|
||
"ftpo_sample_rejected_regularisation_strength",
|
||
"ftpo_sample_chosen_regularisation_strength",
|
||
"ftpo_sample_min_chosen_tokens",
|
||
)
|
||
|
||
|
||
def _deep_update(dst: Dict, src: Dict) -> Dict:
|
||
"""Recursively merge src into dst (src wins)."""
|
||
for k, v in src.items():
|
||
if k in dst and isinstance(v, dict) and isinstance(dst[k], dict):
|
||
_deep_update(dst[k], v)
|
||
else:
|
||
dst[k] = copy.deepcopy(v)
|
||
return dst
|
||
|
||
def load_pipeline_config(config_path: Path) -> Dict[str, Any]:
|
||
"""Load config from a YAML file, or return empty dict if missing/invalid."""
|
||
if config_path and config_path.exists():
|
||
try:
|
||
with config_path.open('r', encoding='utf-8') as f:
|
||
data = yaml.safe_load(f) or {}
|
||
logger.info("Loaded configuration from %s", config_path)
|
||
return data
|
||
except Exception as e:
|
||
logger.error("Could not load %s: %s – using empty config", config_path, e)
|
||
else:
|
||
logger.info("Config file %s not found – using empty config", config_path)
|
||
return {}
|
||
|
||
def merge_config_with_cli_args(config: Dict[str, Any], cli_args: argparse.Namespace) -> Dict[str, Any]:
|
||
"""
|
||
Merges every possible CLI parameter from your old DEFAULT_CONFIG
|
||
into 'config' if the user actually provided it (i.e. it's not None).
|
||
Also merges housekeeping flags (config_file, resume_from_dir, log_level).
|
||
"""
|
||
merged = copy.deepcopy(config)
|
||
|
||
# 1. Housekeeping arguments (not originally in DEFAULT_CONFIG, but we keep them if set)
|
||
if getattr(cli_args, 'config_file', None) is not None:
|
||
merged['config_file'] = cli_args.config_file
|
||
if getattr(cli_args, 'resume_from_dir', None) is not None:
|
||
merged['resume_from_dir'] = cli_args.resume_from_dir
|
||
if getattr(cli_args, 'log_level', None) is not None:
|
||
merged['log_level'] = cli_args.log_level
|
||
|
||
# 2. Booleans that map from CLI flags to known keys in config
|
||
if getattr(cli_args, 'run_finetune', None) is not None:
|
||
merged['finetune_enabled'] = cli_args.run_finetune
|
||
if getattr(cli_args, 'manage_vllm', None) is not None:
|
||
merged['manage_vllm'] = cli_args.manage_vllm
|
||
if getattr(cli_args, 'generation_step_enabled', None) is not None:
|
||
merged['generation_step_enabled'] = cli_args.generation_step_enabled
|
||
if getattr(cli_args, "finetune_cuda_visible_devices", None) is not None:
|
||
merged["finetune_cuda_visible_devices"] = cli_args.finetune_cuda_visible_devices
|
||
|
||
|
||
# 3. All remaining keys from the old DEFAULT_CONFIG
|
||
_all_groups: Sequence[Sequence[str]] = (
|
||
_ALWAYS,
|
||
_VLLM,
|
||
_GENERATION,
|
||
_NGRAM,
|
||
_SLOP,
|
||
_OVERREP,
|
||
_FINETUNE,
|
||
)
|
||
all_config_keys: List[str] = [k for group in _all_groups for k in group]
|
||
|
||
# 4. fallback for per-stage model IDs <-- add this block
|
||
for key in (
|
||
"vllm_model_id",
|
||
"generation_model_id",
|
||
"generation_chat_template_model_id",
|
||
"finetune_base_model_id",
|
||
):
|
||
if not merged.get(key): # None, "", or missing
|
||
merged[key] = merged.get("model_id")
|
||
|
||
# Overwrite config if user specified a value
|
||
for key in all_config_keys:
|
||
cli_val = getattr(cli_args, key, None)
|
||
if cli_val is not None:
|
||
merged[key] = cli_val
|
||
|
||
return merged
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Validate with partial requirements depending on which features are enabled
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _missing(cfg: Dict[str, Any], keys: Sequence[str]) -> List[str]:
|
||
return [k for k in keys if k not in cfg or cfg[k] is None]
|
||
|
||
def validate_config(cfg: Dict[str, Any]) -> None:
|
||
"""Raise ValueError if any required config is missing based on pipeline flags."""
|
||
missing = []
|
||
# always
|
||
missing.extend(_missing(cfg, _ALWAYS))
|
||
|
||
# vllm
|
||
if cfg.get("manage_vllm", False):
|
||
missing.extend(_missing(cfg, _VLLM))
|
||
|
||
# generation
|
||
if cfg.get("generation_step_enabled", True):
|
||
missing.extend(_missing(cfg, _GENERATION))
|
||
|
||
# n-gram ban
|
||
if cfg.get("enable_ngram_ban", False):
|
||
missing.extend(_missing(cfg, _NGRAM))
|
||
|
||
# slop phrase ban
|
||
if cfg.get("enable_slop_phrase_ban", False):
|
||
missing.extend(_missing(cfg, _SLOP))
|
||
|
||
# over-rep analysis
|
||
if cfg.get("compute_overrep_words", False):
|
||
missing.extend(_missing(cfg, _OVERREP))
|
||
|
||
# finetuning
|
||
if cfg.get("finetune_enabled", False):
|
||
missing.extend(_missing(cfg, _FINETUNE))
|
||
|
||
if missing:
|
||
raise ValueError(
|
||
f"Configuration is incomplete; missing these keys: {', '.join(sorted(set(missing)))}"
|
||
)
|
||
logger.info("Configuration validated – all required keys present (for enabled features).")
|
||
|
||
def load_merge_validate(config_path: Path, cli_args: argparse.Namespace) -> Dict[str, Any]:
|
||
"""
|
||
1) Load YAML from config_path,
|
||
2) Merge in any CLI flags user typed,
|
||
3) Validate that all needed keys for enabled features are present.
|
||
"""
|
||
cfg = load_pipeline_config(config_path)
|
||
cfg = merge_config_with_cli_args(cfg, cli_args)
|
||
validate_config(cfg)
|
||
return cfg
|