wehub-resource-sync
49 lines
2.1 KiB
Python
49 lines
2.1 KiB
Python
from pathlib import Path
|
|
|
|
from rdagent.app.finetune.llm.conf import FT_RD_SETTING
|
|
from rdagent.log import rdagent_logger as logger
|
|
from rdagent.scenarios.finetune.datasets import prepare as prepare_dataset
|
|
from rdagent.scenarios.finetune.download.hf import download_model
|
|
|
|
|
|
def ensure_ft_assets_exist(
|
|
*, model: str | None = None, dataset: str | None = None, check_model: bool = False, check_dataset: bool = False
|
|
) -> None:
|
|
"""Ensure dataset and model assets exist under FT_FILE_PATH structure.
|
|
|
|
Args:
|
|
model: Model name to check/download. Required if check_model=True.
|
|
dataset: Dataset name (registered in DATASETS) to check/download. Required if check_dataset=True.
|
|
check_model: Whether to ensure model exists.
|
|
check_dataset: Whether to ensure dataset exists.
|
|
|
|
Paths:
|
|
- Dataset path: FT_RD_SETTING.file_path/datasets/<dataset>
|
|
- Model path: FT_RD_SETTING.file_path/models/<model>
|
|
"""
|
|
# Ensure dataset exists if requested
|
|
if check_dataset:
|
|
if dataset is None:
|
|
raise ValueError("Dataset name is required when check_dataset=True")
|
|
|
|
dataset_dir = Path(FT_RD_SETTING.file_path) / "datasets" / dataset
|
|
if not dataset_dir.exists():
|
|
try:
|
|
logger.info(f"Preparing dataset '{dataset}' to {dataset_dir}")
|
|
prepare_dataset(dataset)
|
|
except Exception as e:
|
|
raise Exception(f"Failed to prepare dataset '{dataset}' to {dataset_dir}: {e}") from e
|
|
|
|
# Ensure model exists if requested
|
|
if check_model:
|
|
if model is None:
|
|
raise ValueError("Model name is required when check_model=True")
|
|
|
|
model_dir = Path(FT_RD_SETTING.file_path) / "models" / model
|
|
if not model_dir.exists():
|
|
try:
|
|
logger.info(f"Downloading model '{model}' to {model_dir}")
|
|
download_model(model, out_dir_root=str(Path(FT_RD_SETTING.file_path) / "models"))
|
|
except Exception as e:
|
|
raise Exception(f"Failed to download model '{model}' to {model_dir}: {e}. ") from e
|