Spaces:
Sleeping
Sleeping
File size: 5,446 Bytes
2d52135 | 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 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | """Helpers for loading model configurations from local trained_models folders."""
from __future__ import annotations
import os
import shutil
from collections.abc import Sequence
from pathlib import Path
from huggingface_hub import snapshot_download
from hydra import compose, initialize_config_dir
from hydra.core.global_hydra import GlobalHydra
from omegaconf import DictConfig
MODELS_ENV = "MINI_TRANSFORMER_MODELS"
MODEL_NAME_ENV = "MINI_TRANSFORMER_MODEL_NAME"
CONFIG_DIR_ENV = "MINI_TRANSFORMER_CONFIG_DIR"
CONFIG_NAME_ENV = "MINI_TRANSFORMER_CONFIG_NAME"
DEFAULT_CONFIG_NAME = "config_inference"
CONFIG_FILENAME = f"{DEFAULT_CONFIG_NAME}.yaml"
def get_models_root() -> Path:
"""Return the directory where local models are stored."""
env_root = os.environ.get(MODELS_ENV)
if env_root:
return Path(env_root).expanduser().resolve()
return (Path.cwd() / "trained_models").resolve()
def ensure_models_root() -> Path:
root = get_models_root()
root.mkdir(parents=True, exist_ok=True)
return root
def list_model_names() -> list[str]:
root = ensure_models_root()
names: list[str] = []
for entry in sorted(root.iterdir()):
if entry.is_dir() and (entry / "configs" / CONFIG_FILENAME).exists():
names.append(entry.name)
return names
def _resolve_relative_paths(cfg: DictConfig, base_dir: Path) -> None:
def _resolve(value: str | None) -> str | None:
if not value:
return value
path = Path(value)
if path.is_absolute():
return str(path)
parts = path.parts
if len(parts) >= 2 and parts[0] == "trained_models" and parts[1] == base_dir.name:
path = Path(*parts[2:])
elif parts and parts[0] == base_dir.name:
path = Path(*parts[1:])
return str((base_dir / path).resolve())
if "model" in cfg:
for key in ("best_checkpoint_path", "latest_checkpoint_path"):
if key in cfg.model:
cfg.model[key] = _resolve(cfg.model[key])
if "tokenizer" in cfg and "path" in cfg.tokenizer:
cfg.tokenizer["path"] = _resolve(cfg.tokenizer["path"])
if "runtime" in cfg:
for key in ("output_dir", "data_dir", "cache_dir", "tokenizer_dir", "checkpoint_path"):
if key in cfg.runtime:
cfg.runtime[key] = _resolve(cfg.runtime[key])
def compose_model_config(
model_name: str,
*,
config_name: str = DEFAULT_CONFIG_NAME,
overrides: Sequence[str] | None = None,
job_name: str | None = None,
) -> DictConfig:
"""Compose the Hydra config for a given model folder."""
models_root = ensure_models_root()
model_dir = models_root / model_name
config_dir = model_dir / "configs"
if not config_dir.is_dir():
raise FileNotFoundError(
f"Config directory not found for model '{model_name}' at {config_dir}."
)
overrides = list(overrides or [])
job = job_name or f"mini_transformer_{model_name}"
if GlobalHydra.instance().is_initialized():
GlobalHydra.instance().clear()
with initialize_config_dir(config_dir=str(config_dir), job_name=job, version_base=None):
cfg = compose(config_name=config_name, overrides=overrides)
_resolve_relative_paths(cfg, model_dir)
return cfg
def compose_config_from_dir(
config_dir: str,
*,
config_name: str = DEFAULT_CONFIG_NAME,
overrides: Sequence[str] | None = None,
job_name: str = "mini_transformer_cli",
) -> DictConfig:
overrides = list(overrides or [])
config_path = Path(config_dir).resolve()
base_dir = config_path.parent if config_path.name == "configs" else config_path
if GlobalHydra.instance().is_initialized():
GlobalHydra.instance().clear()
with initialize_config_dir(config_dir=str(config_path), job_name=job_name, version_base=None):
cfg = compose(config_name=config_name, overrides=overrides)
_resolve_relative_paths(cfg, base_dir)
return cfg
__all__ = [
"CONFIG_FILENAME",
"DEFAULT_CONFIG_NAME",
"MODELS_ENV",
"MODEL_NAME_ENV",
"CONFIG_DIR_ENV",
"CONFIG_NAME_ENV",
"compose_model_config",
"compose_config_from_dir",
"ensure_models_root",
"get_models_root",
"list_model_names",
"snapshot_to_local",
]
def snapshot_to_local(
model_id: str,
*,
local_name: str | None = None,
revision: str | None = None,
repo_type: str = "model",
token: str | None = None,
cache_dir: str | None = None,
force: bool = False,
) -> Path:
"""Download a repository from the Hugging Face Hub into `trained_models/`."""
root = ensure_models_root()
name = local_name or model_id.replace("/", "__")
destination = root / name
if destination.exists() and any(destination.iterdir()):
if not force:
raise FileExistsError(
f"Destination {destination} already exists. Use --force to overwrite."
)
for child in destination.iterdir():
if child.is_dir():
shutil.rmtree(child)
else:
child.unlink()
destination.mkdir(parents=True, exist_ok=True)
snapshot_download(
repo_id=model_id,
repo_type=repo_type,
local_dir=str(destination),
local_dir_use_symlinks=False,
revision=revision,
token=token,
cache_dir=cache_dir,
)
return destination
|