diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/base_model_trpgner/__init__.py | 2 | ||||
| -rw-r--r-- | src/base_model_trpgner/download_model.py | 70 | ||||
| -rw-r--r-- | src/base_model_trpgner/inference/__init__.py | 137 | ||||
| -rw-r--r-- | src/base_model_trpgner/training/__init__.py | 4 |
4 files changed, 122 insertions, 91 deletions
diff --git a/src/base_model_trpgner/__init__.py b/src/base_model_trpgner/__init__.py index 9796c83..8549a69 100644 --- a/src/base_model_trpgner/__init__.py +++ b/src/base_model_trpgner/__init__.py @@ -22,7 +22,7 @@ try: from importlib.metadata import version __version__ = version("base_model_trpgner") except Exception: - __version__ = "0.1.1.dev" + __version__ = "0.1.2.dev" __all__ = [ "__version__", diff --git a/src/base_model_trpgner/download_model.py b/src/base_model_trpgner/download_model.py deleted file mode 100644 index 2d65099..0000000 --- a/src/base_model_trpgner/download_model.py +++ /dev/null @@ -1,70 +0,0 @@ -#!/usr/bin/env python3 -""" -模型下载脚本 - -自动下载预训练的 ONNX 模型到用户缓存目录。 -""" - -import os -import sys -from pathlib import Path -import urllib.request - - -def download_model( - output_dir: str = None, - url: str = "https://github.com/HydroRoll-Team/base-model/releases/download/v0.1.0/model.onnx" -): - """ - 下载 ONNX 模型 - - Args: - output_dir: 输出目录，默认为 ~/.cache/basemodel/models/trpg-final/ - url: 模型下载 URL - """ - if output_dir is None: - output_dir = Path.home() / ".cache" / "basemodel" / "models" / "trpg-final" - else: - output_dir = Path(output_dir) - - output_dir.mkdir(parents=True, exist_ok=True) - output_path = output_dir / "model.onnx" - - if output_path.exists(): - print(f"✅ 模型已存在: {output_path}") - return str(output_path) - - print(f"📥 正在下载模型到 {output_path}...") - print(f" URL: {url}") - - try: - urllib.request.urlretrieve(url, output_path) - print(f"✅ 模型下载成功!") - return str(output_path) - except Exception as e: - print(f"❌ 下载失败: {e}") - print(f" 请手动从以下地址下载模型:") - print(f" {url}") - sys.exit(1) - - -if __name__ == "__main__": - import argparse - - parser = argparse.ArgumentParser(description="下载 base-model ONNX 模型") - parser.add_argument( - "--output-dir", - type=str, - default=None, - help="模型输出目录（默认: ~/.cache/basemodel/models/trpg-final/）" - ) - parser.add_argument( - "--url", - type=str, - default="https://github.com/HydroRoll-Team/base-model/releases/download/v0.1.0/model.onnx", - help="模型下载 URL" - ) - - args = parser.parse_args() - - download_model(args.output_dir, args.url) diff --git a/src/base_model_trpgner/inference/__init__.py b/src/base_model_trpgner/inference/__init__.py index 93a185f..d70cb23 100644 --- a/src/base_model_trpgner/inference/__init__.py +++ b/src/base_model_trpgner/inference/__init__.py @@ -5,6 +5,8 @@ ONNX 推理模块 """ import os +import json +import shutil from typing import List, Dict, Any, Optional from pathlib import Path @@ -18,20 +20,113 @@ except ImportError as e: ) from e -# 默认模型路径（相对于包安装位置） -DEFAULT_MODEL_DIR = Path(__file__).parent.parent.parent.parent / "models" / "trpg-final" -# 远程模型 URL（用于自动下载） -MODEL_URL = "https://github.com/HydroRoll-Team/base-model/releases/download/v0.1.0/model.onnx" +# GitHub 仓库信息 +REPO_OWNER = "HydroRoll-Team" +REPO_NAME = "base-model" +# 用户数据目录 +USER_MODEL_DIR = Path.home() / ".cache" / "base_model_trpgner" / "models" / "trpg-final" + + +def get_latest_release_url() -> str: + """ + 获取 GitHub 最新 Release 的下载 URL + + Returns: + 最新 Release 的标签名（如 v0.1.0） + """ + import urllib.request + import urllib.error + + try: + # 使用 GitHub API 获取最新 release + api_url = f"https://api.github.com/repos/{REPO_OWNER}/{REPO_NAME}/releases/latest" + with urllib.request.urlopen(api_url, timeout=10) as response: + data = json.load(response) + return data.get("tag_name", "v0.1.0") + except (urllib.error.URLError, json.JSONDecodeError, KeyError): + # 失败时返回默认版本 + return "v0.1.0" + + +def download_model_files(version: Optional[str] = None, force: bool = False) -> Path: + """ + 从 GitHub Release 下载模型文件 + + Args: + version: Release 版本（如 v0.1.0），None 表示最新版本 + force: 是否强制重新下载（即使文件已存在） + + Returns: + 模型文件保存目录 + """ + import urllib.request + import urllib.error + + if version is None: + version = get_latest_release_url() + + model_dir = USER_MODEL_DIR + model_dir.mkdir(parents=True, exist_ok=True) + + # 检查是否已下载 + marker_file = model_dir / ".version" + if not force and marker_file.exists(): + with open(marker_file, "r") as f: + current_version = f.read().strip() + if current_version == version: + print(f"模型已存在 (版本: {version})") + return model_dir + + print(f"正在下载模型 {version}...") + + # 需要下载的文件 + base_url = f"https://github.com/{REPO_OWNER}/{REPO_NAME}/releases/download/{version}" + files_to_download = [ + "model.onnx", + "model.onnx.data", + "config.json", + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", + "vocab.txt", + ] + + for filename in files_to_download: + url = f"{base_url}/{filename}" + dest_path = model_dir / filename + + try: + print(f"下载 {filename}...") + with urllib.request.urlopen(url, timeout=60) as response: + with open(dest_path, "wb") as f: + shutil.copyfileobj(response, f) + except urllib.error.HTTPError as e: + print(f"下载 {filename} 失败: {e}") + # model.onnx.data 不是必需的（某些小模型可能没有） + if filename == "model.onnx.data": + continue + else: + raise RuntimeError(f"无法下载模型文件: {filename}") from e + + # 写入版本标记 + with open(marker_file, "w") as f: + f.write(version) + + print(f"✅ 模型下载完成: {model_dir}") + return model_dir class TRPGParser: """ TRPG 日志解析器（基于 ONNX） + 首次运行时会自动从 GitHub Release 下载最新模型。 + Args: - model_path: ONNX 模型路径，默认使用内置模型 + model_path: ONNX 模型路径，默认使用自动下载的模型 tokenizer_path: tokenizer 配置路径，默认与 model_path 相同 device: 推理设备，"cpu" 或 "cuda" + auto_download: 是否自动下载模型（默认 True） Examples: >>> parser = TRPGParser() @@ -45,10 +140,11 @@ class TRPGParser: model_path: Optional[str] = None, tokenizer_path: Optional[str] = None, device: str = "cpu", + auto_download: bool = True, ): # 确定模型路径 if model_path is None: - model_path = self._get_default_model_path() + model_path = self._get_default_model_path(auto_download) if tokenizer_path is None: tokenizer_path = Path(model_path).parent @@ -60,25 +156,31 @@ class TRPGParser: # 加载模型 self._load_model() - def _get_default_model_path(self) -> str: - """获取默认模型路径""" - # 1. 尝试相对于项目根目录 + def _get_default_model_path(self, auto_download: bool) -> str: + """获取默认模型路径，必要时自动下载""" + # 1. 检查本地开发环境 project_root = Path(__file__).parent.parent.parent.parent local_model = project_root / "models" / "trpg-final" / "model.onnx" if local_model.exists(): return str(local_model) - # 2. 尝试用户数据目录 - from pathlib import Path - user_model_dir = Path.home() / ".cache" / "basemodel" / "models" / "trpg-final" - user_model = user_model_dir / "model.onnx" + # 2. 检查用户缓存目录 + user_model = USER_MODEL_DIR / "model.onnx" if user_model.exists(): return str(user_model) - # 3. 抛出错误，提示下载 + # 3. 自动下载 + if auto_download: + print("模型未找到，正在从 GitHub Release 下载...") + download_model_files() + return str(user_model) + + # 4. 抛出错误 raise FileNotFoundError( - f"模型文件未找到。请从 {MODEL_URL} 下载模型到 {user_model_dir}\n" - f"或运行: python -m basemodel.download_model" + f"模型文件未找到。\n" + f"请开启自动下载: TRPGParser(auto_download=True)\n" + f"或手动下载到: {USER_MODEL_DIR}\n" + f"下载地址: https://github.com/{REPO_OWNER}/{REPO_NAME}/releases/latest" ) def _load_model(self): @@ -100,7 +202,6 @@ class TRPGParser: ) # 加载标签映射 - import json config_path = self.tokenizer_path / "config.json" if config_path.exists(): with open(config_path, "r", encoding="utf-8") as f: @@ -289,4 +390,4 @@ def parse_lines(texts: List[str], model_path: Optional[str] = None) -> List[Dict return parser.parse_batch(texts) -__all__ = ["TRPGParser", "parse_line", "parse_lines"] +__all__ = ["TRPGParser", "parse_line", "parse_lines", "download_model_files"] diff --git a/src/base_model_trpgner/training/__init__.py b/src/base_model_trpgner/training/__init__.py index ccf3c03..4f8e30d 100644 --- a/src/base_model_trpgner/training/__init__.py +++ b/src/base_model_trpgner/training/__init__.py @@ -36,11 +36,11 @@ def train_ner_model( resume_from_checkpoint: 恢复检查点路径 Examples: - >>> from basemodeltrpgner.training import train_ner_model + >>> from base_model_trpgner.training import train_ner_model >>> train_ner_model( ... conll_data="./data", ... output_dir="./my_model", - ... epochs=10 + ... num_train_epochs=10 ... ) """ try: |