| |
| |
| """ |
| 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) |
|
|
| |
| 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 |
| |
| 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" |
|
|
