File size: 4,550 Bytes
3fdd0bf | 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 | #!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
HOS Model Optimizer - 工具函数模块
提供通用工具函数：
- 日志配置
- 文件操作工具
- 模型路径处理
"""
import os
import sys
import logging
from pathlib import Path
from typing import Optional, List
# ============================================================
# 日志配置
# ============================================================
def setup_logger(
name: str = "hos_optimizer",
level: int = logging.INFO,
log_file: Optional[str] = None,
fmt: str = "%(asctime)s [%(levelname)s] %(name)s: %(message)s",
datefmt: str = "%H:%M:%S",
) -> logging.Logger:
"""
配置并返回日志记录器
Args:
name: 日志记录器名称
level: 日志级别
log_file: 日志文件路径（可选）
fmt: 日志格式
datefmt: 日期格式
Returns:
配置好的 Logger 实例
"""
logger = logging.getLogger(name)
logger.setLevel(level)
# 避免重复添加 handler
if logger.handlers:
return logger
formatter = logging.Formatter(fmt, datefmt=datefmt)
# 控制台输出
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setFormatter(formatter)
logger.addHandler(console_handler)
# 文件输出（可选）
if log_file:
file_handler = logging.FileHandler(log_file, encoding="utf-8")
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
return logger
# ============================================================
# 文件操作工具
# ============================================================
def ensure_dir(path: str) -> str:
"""
确保目录存在，不存在则创建
Args:
path: 目录路径
Returns:
目录路径
"""
Path(path).mkdir(parents=True, exist_ok=True)
return path
def get_file_size_gb(path: str) -> float:
"""
获取文件大小（GB）
Args:
path: 文件路径
Returns:
文件大小（GB）
"""
return os.path.getsize(path) / (1024 ** 3)
def get_dir_size_gb(path: str) -> float:
"""
获取目录总大小（GB）
Args:
path: 目录路径
Returns:
目录总大小（GB）
"""
total = 0
for dirpath, _, filenames in os.walk(path):
for f in filenames:
fp = os.path.join(dirpath, f)
if os.path.isfile(fp):
total += os.path.getsize(fp)
return total / (1024 ** 3)
def find_model_files(path: str) -> List[str]:
"""
在目录中查找模型文件
Args:
path: 搜索路径
Returns:
模型文件路径列表
"""
extensions = (".safetensors", ".bin", ".pt", ".gguf", ".onnx")
model_files = []
for dirpath, _, filenames in os.walk(path):
for f in filenames:
if f.endswith(extensions):
model_files.append(os.path.join(dirpath, f))
return sorted(model_files)
# ============================================================
# 模型路径处理
# ============================================================
def resolve_model_path(path: str) -> str:
"""
解析模型路径，支持相对路径和环境变量展开
Args:
path: 原始路径
Returns:
解析后的绝对路径
"""
expanded = os.path.expandvars(os.path.expanduser(path))
return os.path.abspath(expanded)
def is_model_path(path: str) -> bool:
"""
判断路径是否为有效的模型路径（本地目录或 HF Hub ID）
Args:
path: 路径字符串
Returns:
是否为有效模型路径
"""
# 本地路径检查
if os.path.exists(path):
return True
# HF Hub ID 格式检查（如 "Qwen/Qwen2.5-0.5B"）
if "/" in path and not os.path.sep in path.replace("/", os.path.sep):
parts = path.split("/")
if len(parts) == 2 and all(parts):
return True
return False
def get_model_format(path: str) -> str:
"""
推断模型格式
Args:
path: 模型路径
Returns:
格式字符串：gguf / safetensors / pytorch / unknown
"""
if path.endswith(".gguf"):
return "gguf"
if os.path.isdir(path):
files = find_model_files(path)
for f in files:
if f.endswith(".safetensors"):
return "safetensors"
if f.endswith((".bin", ".pt")):
return "pytorch"
return "unknown"
|