File size: 6,761 Bytes
bcc2f7b |
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 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 |
"""
Model utilities for downloading and managing the marine species model.
"""
import os
import shutil
from pathlib import Path
from typing import Optional, Dict, Any
from huggingface_hub import hf_hub_download, list_repo_files
from app.core.config import settings
from app.core.logging import get_logger
logger = get_logger(__name__)
def download_model_from_hf(
repo_id: str,
model_filename: str,
local_dir: str,
force_download: bool = False
) -> str:
"""
Download model from HuggingFace Hub.
Args:
repo_id: HuggingFace repository ID
model_filename: Name of the model file
local_dir: Local directory to save the model
force_download: Whether to force re-download if file exists
Returns:
Path to the downloaded model file
"""
try:
# Create local directory if it doesn't exist
Path(local_dir).mkdir(parents=True, exist_ok=True)
local_path = Path(local_dir) / model_filename
# Check if file already exists and force_download is False
if local_path.exists() and not force_download:
logger.info(f"Model already exists at {local_path}")
return str(local_path)
logger.info(f"Downloading {model_filename} from {repo_id}...")
downloaded_path = hf_hub_download(
repo_id=repo_id,
filename=model_filename,
local_dir=local_dir,
local_dir_use_symlinks=False,
force_download=force_download
)
logger.info(f"Model downloaded successfully to: {downloaded_path}")
return downloaded_path
except Exception as e:
logger.error(f"Failed to download model: {str(e)}")
raise
def list_available_files(repo_id: str) -> list:
"""
List all available files in a HuggingFace repository.
Args:
repo_id: HuggingFace repository ID
Returns:
List of available files
"""
try:
files = list_repo_files(repo_id)
return files
except Exception as e:
logger.error(f"Failed to list repository files: {str(e)}")
return []
def verify_model_file(model_path: str) -> bool:
"""
Verify that a model file exists and is valid.
Args:
model_path: Path to the model file
Returns:
True if model file is valid
"""
try:
path = Path(model_path)
# Check if file exists
if not path.exists():
logger.error(f"Model file does not exist: {model_path}")
return False
# Check file size (should be > 1MB for a real model)
file_size = path.stat().st_size
if file_size < 1024 * 1024: # 1MB
logger.warning(f"Model file seems too small: {file_size} bytes")
return False
# Check file extension
if not path.suffix.lower() in ['.pt', '.pth']:
logger.warning(f"Unexpected model file extension: {path.suffix}")
logger.info(f"Model file verified: {model_path} ({file_size / (1024*1024):.1f} MB)")
return True
except Exception as e:
logger.error(f"Failed to verify model file: {str(e)}")
return False
def get_model_info(model_path: str) -> Dict[str, Any]:
"""
Get information about a model file.
Args:
model_path: Path to the model file
Returns:
Dictionary with model information
"""
info = {
"path": model_path,
"exists": False,
"size_mb": 0,
"size_bytes": 0
}
try:
path = Path(model_path)
if path.exists():
info["exists"] = True
size_bytes = path.stat().st_size
info["size_bytes"] = size_bytes
info["size_mb"] = size_bytes / (1024 * 1024)
info["modified_time"] = path.stat().st_mtime
except Exception as e:
logger.error(f"Failed to get model info: {str(e)}")
return info
def cleanup_model_cache(cache_dir: Optional[str] = None) -> None:
"""
Clean up model cache directory.
Args:
cache_dir: Cache directory to clean (uses default if None)
"""
try:
if cache_dir is None:
cache_dir = Path.home() / ".cache" / "huggingface"
cache_path = Path(cache_dir)
if cache_path.exists():
logger.info(f"Cleaning up cache directory: {cache_path}")
shutil.rmtree(cache_path)
logger.info("Cache cleanup completed")
else:
logger.info("Cache directory does not exist")
except Exception as e:
logger.error(f"Failed to cleanup cache: {str(e)}")
def setup_model_directory() -> str:
"""
Setup the model directory and ensure it exists.
Returns:
Path to the model directory
"""
model_dir = Path(settings.MODEL_PATH).parent
model_dir.mkdir(parents=True, exist_ok=True)
logger.info(f"Model directory setup: {model_dir}")
return str(model_dir)
if __name__ == "__main__":
# Command line utility for model management
import argparse
parser = argparse.ArgumentParser(description="Model management utility")
parser.add_argument("--download", action="store_true", help="Download model from HuggingFace")
parser.add_argument("--verify", action="store_true", help="Verify model file")
parser.add_argument("--info", action="store_true", help="Show model information")
parser.add_argument("--list-files", action="store_true", help="List available files in HF repo")
parser.add_argument("--cleanup-cache", action="store_true", help="Cleanup model cache")
parser.add_argument("--force", action="store_true", help="Force download even if file exists")
args = parser.parse_args()
if args.download:
setup_model_directory()
download_model_from_hf(
repo_id=settings.HUGGINGFACE_REPO,
model_filename=f"{settings.MODEL_NAME}.pt",
local_dir=str(Path(settings.MODEL_PATH).parent),
force_download=args.force
)
if args.verify:
is_valid = verify_model_file(settings.MODEL_PATH)
print(f"Model valid: {is_valid}")
if args.info:
info = get_model_info(settings.MODEL_PATH)
print(f"Model info: {info}")
if args.list_files:
files = list_available_files(settings.HUGGINGFACE_REPO)
print(f"Available files in {settings.HUGGINGFACE_REPO}:")
for file in files:
print(f" - {file}")
if args.cleanup_cache:
cleanup_model_cache()
|