datawhalechina--self-llm
wehub-resource-sync
251 行
8.3 KiB
Python
251 行
8.3 KiB
Python
"""
|
||
模型下载模块
|
||
从 configs/model_info/ 目录读取模型列表，支持 mlx / original 两种来源
|
||
模型按 框架/Company/Series/ModelName 目录结构存放
|
||
|
||
独立运行: python -m modules.download_model
|
||
"""
|
||
|
||
import os
|
||
import json
|
||
import time
|
||
from huggingface_hub import snapshot_download
|
||
|
||
# ============================================================
|
||
# 📦 配置
|
||
# ============================================================
|
||
|
||
ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
LOCAL_DIR = os.path.join(ROOT_DIR, "models")
|
||
CONFIGS_DIR = os.path.join(ROOT_DIR, "configs", "model_info")
|
||
|
||
|
||
# ============================================================
|
||
# 🔧 核心函数
|
||
# ============================================================
|
||
|
||
def load_models_config(source="mlx"):
|
||
"""
|
||
从 JSON 配置文件加载模型列表
|
||
|
||
Returns:
|
||
list[dict]: [{"Company": "...", "Series": "...", "Models": [...]}]
|
||
"""
|
||
filename = f"{source}.json"
|
||
config_path = os.path.join(CONFIGS_DIR, filename)
|
||
if not os.path.exists(config_path):
|
||
return []
|
||
with open(config_path, "r", encoding="utf-8") as f:
|
||
return json.load(f)
|
||
|
||
|
||
def get_companies(source="mlx"):
|
||
"""获取所有公司名称"""
|
||
config = load_models_config(source)
|
||
return sorted(set(item["Company"] for item in config))
|
||
|
||
|
||
def get_series(company, source="mlx"):
|
||
"""获取指定公司的所有系列"""
|
||
config = load_models_config(source)
|
||
return sorted(set(
|
||
item["Series"] for item in config if item["Company"] == company
|
||
))
|
||
|
||
|
||
def get_models(company, series, source="mlx"):
|
||
"""获取指定公司 + 系列下的模型列表"""
|
||
config = load_models_config(source)
|
||
for item in config:
|
||
if item["Company"] == company and item["Series"] == series:
|
||
return item.get("Models", [])
|
||
return []
|
||
|
||
|
||
def get_framework_inference(company, series, source="mlx"):
|
||
"""获取指定公司+系列支持的推理框架列表"""
|
||
from modules.core_types import Framework
|
||
config = load_models_config(source)
|
||
for item in config:
|
||
if item["Company"] == company and item["Series"] == series:
|
||
frameworks = item.get("FrameworkInference")
|
||
if frameworks:
|
||
return [Framework(f) for f in frameworks]
|
||
# 未配置时按来源默认
|
||
if source == "mlx":
|
||
return [Framework.MLX]
|
||
return [Framework.TRANSFORMERS]
|
||
|
||
|
||
def find_model_info(model_name, source="mlx"):
|
||
"""根据模型名称反查 Company 和 Series"""
|
||
config = load_models_config(source)
|
||
for item in config:
|
||
if model_name in item.get("Models", []):
|
||
return item["Company"], item["Series"]
|
||
return None, None
|
||
|
||
|
||
def get_repo_id(model_name, source="mlx"):
|
||
"""根据模型名称和来源拼接 repo_id"""
|
||
if source == "mlx":
|
||
return f"mlx-community/{model_name}"
|
||
return model_name
|
||
|
||
|
||
def get_local_path(model_name, source="mlx", company=None, series=None):
|
||
"""
|
||
获取模型本地路径: models/source/Company/Series/ModelName
|
||
|
||
company/series 可选，未提供时自动从配置中查找
|
||
"""
|
||
if not company or not series:
|
||
company, series = find_model_info(model_name, source)
|
||
|
||
repo_id = get_repo_id(model_name, source)
|
||
local_name = repo_id.split("/")[-1]
|
||
|
||
if company and series:
|
||
return os.path.join(LOCAL_DIR, source, company, series, local_name)
|
||
# 兜底：找不到配置时放在 models/source/ 下
|
||
return os.path.join(LOCAL_DIR, source, local_name)
|
||
|
||
|
||
def model_exists(model_name, source="mlx", company=None, series=None):
|
||
"""检测本地是否已存在该模型"""
|
||
local_path = get_local_path(model_name, source, company, series)
|
||
return os.path.exists(os.path.join(local_path, "config.json"))
|
||
|
||
|
||
def scan_local_models():
|
||
"""
|
||
扫描本地已下载的模型（遍历 source/Company/Series/Model 目录结构）
|
||
|
||
Returns:
|
||
list[dict]: [{"source": "...", "company": "...", "series": "...", "model": "...", "path": "...", "label": "..."}]
|
||
"""
|
||
if not os.path.exists(LOCAL_DIR):
|
||
return []
|
||
results = []
|
||
for source in sorted(os.listdir(LOCAL_DIR)):
|
||
source_dir = os.path.join(LOCAL_DIR, source)
|
||
if not os.path.isdir(source_dir):
|
||
continue
|
||
for company in sorted(os.listdir(source_dir)):
|
||
company_dir = os.path.join(source_dir, company)
|
||
if not os.path.isdir(company_dir):
|
||
continue
|
||
for series in sorted(os.listdir(company_dir)):
|
||
series_dir = os.path.join(company_dir, series)
|
||
if not os.path.isdir(series_dir):
|
||
continue
|
||
for model in sorted(os.listdir(series_dir)):
|
||
model_dir = os.path.join(series_dir, model)
|
||
if os.path.isdir(model_dir) and os.path.exists(os.path.join(model_dir, "config.json")):
|
||
results.append({
|
||
"source": source,
|
||
"company": company,
|
||
"series": series,
|
||
"model": model,
|
||
"path": model_dir,
|
||
"label": f"[{source}] {company} / {series} / {model}",
|
||
})
|
||
return results
|
||
|
||
|
||
def download(model_name, source="mlx", company=None, series=None):
|
||
"""
|
||
下载模型到本地
|
||
|
||
Returns:
|
||
(local_path, elapsed) 下载成功
|
||
Raises:
|
||
FileExistsError: 模型已存在
|
||
"""
|
||
repo_id = get_repo_id(model_name, source)
|
||
local_path = get_local_path(model_name, source, company, series)
|
||
|
||
if os.path.exists(os.path.join(local_path, "config.json")):
|
||
raise FileExistsError(f"模型已存在: {local_path}")
|
||
|
||
print(f"⏳ 开始下载 {repo_id}")
|
||
print(f"📂 保存路径: {local_path}")
|
||
|
||
os.makedirs(local_path, exist_ok=True)
|
||
s = time.time()
|
||
snapshot_download(repo_id=repo_id, local_dir=local_path)
|
||
elapsed = time.time() - s
|
||
|
||
print(f"✅ 下载完成，耗时 {elapsed:.2f} 秒")
|
||
return local_path, elapsed
|
||
|
||
|
||
# ============================================================
|
||
# 🚀 独立运行：交互式下载
|
||
# ============================================================
|
||
|
||
if __name__ == "__main__":
|
||
print("=" * 50)
|
||
print("📥 模型下载工具")
|
||
print("=" * 50)
|
||
|
||
# 选择来源
|
||
print("\n模型来源:")
|
||
print(" 1. mlx（已量化 MLX 格式，推荐 Mac）")
|
||
print(" 2. original（原始 HuggingFace 模型）")
|
||
source_input = input("\n请选择 [1/2]（默认 1）: ").strip() or "1"
|
||
source = "mlx" if source_input == "1" else "original"
|
||
|
||
# 选择公司
|
||
companies = get_companies(source)
|
||
if not companies:
|
||
print("❌ 未找到模型配置，请检查 configs/ 目录")
|
||
exit(1)
|
||
|
||
print(f"\n公司/组织:")
|
||
for i, c in enumerate(companies, 1):
|
||
print(f" {i}. {c}")
|
||
idx = int(input(f"\n请选择 [1-{len(companies)}]（默认 1）: ").strip() or "1") - 1
|
||
company = companies[idx]
|
||
|
||
# 选择系列
|
||
series_list = get_series(company, source)
|
||
print(f"\n模型系列（{company}）:")
|
||
for i, s in enumerate(series_list, 1):
|
||
print(f" {i}. {s}")
|
||
idx = int(input(f"\n请选择 [1-{len(series_list)}]（默认 1）: ").strip() or "1") - 1
|
||
series = series_list[idx]
|
||
|
||
# 选择模型
|
||
models = get_models(company, series, source)
|
||
if not models:
|
||
print("❌ 该系列下暂无模型")
|
||
exit(1)
|
||
|
||
print(f"\n可用模型（{company} / {series}）:")
|
||
for i, name in enumerate(models, 1):
|
||
exists = "✅ 已下载" if model_exists(name, source, company, series) else ""
|
||
print(f" {i}. {name} {exists}")
|
||
|
||
idx = int(input(f"\n请选择模型编号 [1-{len(models)}]（默认 1）: ").strip() or "1") - 1
|
||
model_name = models[idx]
|
||
|
||
# 检测是否存在
|
||
if model_exists(model_name, source, company, series):
|
||
print(f"\n✅ 模型已存在本地，无需下载")
|
||
print(f"📂 路径: {get_local_path(model_name, source, company, series)}")
|
||
exit(0)
|
||
|
||
# 确认下载
|
||
repo_id = get_repo_id(model_name, source)
|
||
confirm = input(f"\n确认下载 {repo_id}？[y/N]: ").strip().lower()
|
||
if confirm != "y":
|
||
print("已取消")
|
||
exit(0)
|
||
|
||
try:
|
||
download(model_name, source, company, series)
|
||
except Exception as e:
|
||
print(f"❌ 下载失败: {e}")
|
||
exit(1)
|