""" 霜云(Shimokumo) - 主入口文件 AI全功能ISP网络运营商巫女模型的启动入口。 支持命令行参数: --host 服务监听地址 (默认: 0.0.0.0) --port 服务监听端口 (默认: 8080) --api-key API认证密钥 --model-path 模型权重文件路径 --tokenizer-path 分词器文件路径 --device 运行设备 (auto/cuda/cpu) --log-level 日志级别 (DEBUG/INFO/WARNING/ERROR) --log-file 日志文件路径 --no-search 禁用搜索功能 --no-novel 禁用小说功能 --no-video 禁用视频功能 --no-chat 禁用角色聊天功能 --no-media 禁用媒体功能 --no-browser 禁用浏览功能 --debug 启用调试模式 --interactive 启动交互式命令行模式(不启动服务器) --help 显示帮助信息 用法: # 启动服务器 python main.py --port 8080 # 交互模式 python main.py --interactive # 加载模型并启动 python main.py --model-path ./models/shimokumo.bin --tokenizer-path ./models/tokenizer.model """ import argparse import os import signal import sys import time # 将src目录加入Python路径 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from config import ShimokumoConfig from utils.logger import get_logger from utils.helpers import format_duration logger = get_logger("Shimokumo.Main") def parse_args() -> argparse.Namespace: """解析命令行参数""" parser = argparse.ArgumentParser( description="霜云(Shimokumo) - AI全功能ISP网络运营商巫女模型", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=( "示例:\n" " python main.py --port 8080 # 启动API服务器\n" " python main.py --interactive # 启动交互式聊天\n" " python main.py --model-path model.bin # 加载模型权重\n" " python main.py --debug --log-level DEBUG # 调试模式\n" ), ) # 服务配置 parser.add_argument("--host", type=str, default="0.0.0.0", help="服务监听地址") parser.add_argument("--port", type=int, default=8080, help="服务监听端口") parser.add_argument("--api-key", type=str, default="", help="API认证密钥(为空则不需要认证)") # 模型配置 parser.add_argument("--model-path", type=str, default="", help="模型权重文件路径") parser.add_argument("--tokenizer-path", type=str, default="", help="分词器模型文件路径") parser.add_argument("--device", type=str, default="auto", choices=["auto", "cuda", "cpu"], help="运行设备") # 日志配置 parser.add_argument("--log-level", type=str, default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR"], help="日志级别") parser.add_argument("--log-file", type=str, default="", help="日志文件路径") # 功能开关 parser.add_argument("--no-search", action="store_true", help="禁用搜索功能") parser.add_argument("--no-novel", action="store_true", help="禁用小说创作功能") parser.add_argument("--no-video", action="store_true", help="禁用视频制作功能") parser.add_argument("--no-chat", action="store_true", help="禁用角色聊天功能") parser.add_argument("--no-media", action="store_true", help="禁用媒体播放功能") parser.add_argument("--no-browser", action="store_true", help="禁用网页浏览功能") # 模式选择 parser.add_argument("--debug", action="store_true", help="启用调试模式") parser.add_argument("--interactive", action="store_true", help="启动交互式命令行聊天模式") return parser.parse_args() def load_config(args: argparse.Namespace) -> ShimokumoConfig: """从命令行参数创建配置""" config = ShimokokuConfig( host=args.host, port=args.port, api_key=args.api_key, device=args.device, model_path=args.model_path, tokenizer_path=args.tokenizer_path, log_level=args.log_level, debug=args.debug, ) # 设置日志文件 if args.log_file: config.log_file = args.log_file if args.debug and not config.log_file: config.log_file = "/data/user/work/shimokumo_source/logs/shimokumo_debug.log" # 功能开关 if args.no_search: config.enable_web_search = False if args.no_novel: config.enable_novel_creation = False if args.no_video: config.enable_video_generation = False if args.no_chat: config.enable_character_chat = False if args.no_media: config.enable_media_player = False if args.no_browser: config.enable_web_browser = False return config def load_model(config: ShimokumoConfig): """ 加载模型和分词器。 Args: config: 模型配置 Returns: (model, tokenizer) 或 (None, None) """ try: import torch from model.shimokumo_model import ShimokumoModel from model.tokenizer import ShimokumoTokenizer # 加载分词器 tokenizer_path = config.tokenizer_path or "" tokenizer = ShimokumoTokenizer(model_path=tokenizer_path if tokenizer_path else None) logger.info(f"分词器加载完成: {tokenizer}") # 加载模型 if config.model_path and os.path.exists(config.model_path): logger.info(f"正在加载模型权重: {config.model_path}") model = ShimokumoModel.from_pretrained(config, config.model_path) logger.info(f"模型加载完成: {model.get_num_params() / 1e9:.2f}B 参数") else: logger.warning("未指定模型权重路径,将使用未训练的随机初始化模型") logger.info("如需加载预训练权重,请使用 --model-path 参数") model = ShimokumoModel(config) logger.info(f"已创建模型架构: {model.get_num_params() / 1e9:.2f}B 参数") return model, tokenizer except Exception as e: logger.error(f"模型加载失败: {e}") return None, None def init_modules(config: ShimokokuConfig, inference=None): """ 初始化功能模块。 Args: config: 模型配置 inference: 推理引擎(可为None) Returns: 模块字典 """ modules: Dict = {} # 联网搜索模块 if config.enable_web_search: try: from modules.web_search import WebSearchModule modules["web_search"] = WebSearchModule() logger.info("联网搜索模块已加载") except Exception as e: logger.warning(f"搜索模块加载失败: {e}") config.enable_web_search = False # 网页浏览模块 if config.enable_web_browser: try: from modules.web_browser import WebBrowserModule modules["web_browser"] = WebBrowserModule() logger.info("网页浏览模块已加载") except Exception as e: logger.warning(f"浏览模块加载失败: {e}") config.enable_web_browser = False # 多媒体播放模块 if config.enable_media_player: try: from modules.media_player import MediaPlayerModule modules["media_player"] = MediaPlayerModule() logger.info("多媒体播放模块已加载") except Exception as e: logger.warning(f"媒体模块加载失败: {e}") config.enable_media_player = False # 小说创作模块 if config.enable_novel_creation: try: from modules.novel_creator import NovelCreatorModule modules["novel_creator"] = NovelCreatorModule(inference_engine=inference) logger.info("小说创作模块已加载") except Exception as e: logger.warning(f"小说模块加载失败: {e}") config.enable_novel_creation = False # 视频制作模块 if config.enable_video_generation: try: from modules.video_generator import VideoGeneratorModule modules["video_generator"] = VideoGeneratorModule(inference_engine=inference) logger.info("视频制作模块已加载") except Exception as e: logger.warning(f"视频模块加载失败: {e}") config.enable_video_generation = False # 角色聊天模块 if config.enable_character_chat: try: from modules.character_chat import CharacterChatModule modules["character_chat"] = CharacterChatModule(inference_engine=inference) logger.info("角色聊天模块已加载") except Exception as e: logger.warning(f"角色聊天模块加载失败: {e}") config.enable_character_chat = False return modules def interactive_mode(config: ShimokokuConfig, inference=None, modules: Dict = None): """ 交互式命令行聊天模式。 Args: config: 模型配置 inference: 推理引擎 modules: 功能模块字典 """ print() print("=" * 50) print(f" 霜云(Shimokumo) - 交互式聊天模式") print(f" {config.character_title}") print("=" * 50) print() print("命令:") print(" /chat <消息> - 角色扮演聊天") print(" /search <关键词> - 联网搜索") print(" /browse - 浏览网页") print(" /clear - 清空对话历史") print(" /info - 查看信息") print(" /quit - 退出") print() print("直接输入消息即可开始普通聊天。") print() # 初始化角色聊天模块 from modules.character_chat import CharacterChatModule chat_module = modules.get("character_chat") or CharacterChatModule(inference_engine=inference) while True: try: user_input = input("\n你: ").strip() except (KeyboardInterrupt, EOFError): print("\n\n再见~!") break if not user_input: continue # 命令处理 if user_input.startswith("/"): parts = user_input.split(maxsplit=1) cmd = parts[0].lower() arg = parts[1] if len(parts) > 1 else "" if cmd == "/quit" or cmd == "/exit": print(f"\n{config.character_name}: 再见~期待下次见到主人的说...\n") break elif cmd == "/clear": if inference: inference.clear_dialogue() chat_module.clear_history() print("对话历史已清空。") elif cmd == "/info": print(f"\n--- 霜云信息 ---") print(f"名称: {config.character_name} ({config.character_name_en})") print(f"头衔: {config.character_title}") print(f"设备: {inference.device if inference else 'N/A'}") if inference: print(f"参数量: {inference.model.get_num_params() / 1e9:.2f}B") print(f"功能: {'、'.join(k.replace('_', '') for k, v in modules.items())}") print() elif cmd == "/search" and arg: if modules.get("web_search"): results = modules["web_search"].search(arg) print(f"\n{modules['web_search'].format_results(results)}") else: print("搜索功能未启用。") elif cmd == "/browse" and arg: if modules.get("web_browser"): page = modules["web_browser"].fetch(arg) if page: summary = modules["web_browser"].generate_summary(page) print(f"\n标题: {page.title}") print(f"摘要: {summary}") else: print("无法获取网页。") else: print("浏览功能未启用。") elif cmd == "/chat" and arg: response = chat_module.chat(arg) print(f"\n{config.character_name}: {response.text}") else: print(f"未知命令: {cmd}。输入 /quit 退出。") else: # 普通聊天 print(f"\n{config.character_name}: ", end="", flush=True) if inference: response = inference.chat_reply(user_input) print(response) else: # 回退:使用角色聊天模块 response = chat_module.chat(user_input) print(response.text) def main(): """主入口函数""" start_time = time.time() # 打印启动横幅 print() print(" ╔══════════════════════════════════════════════════╗") print(" ║ ║") print(" ║ 霜云 (Shimokumo) ║") print(" ║ AI全功能ISP网络运营商巫女模型 ║") print(" ║ ║") print(" ╚══════════════════════════════════════════════════╝") print() # 解析参数 args = parse_args() # 创建配置 config = load_config(args) # 初始化日志 logger = get_logger( "Shimokumo", log_level=config.log_level, log_file=config.log_file if config.log_file else None, ) logger.info("霜云启动中...") logger.info(f"配置: host={config.host}, port={config.port}, device={config.device}") # 加载模型 model, tokenizer = load_model(config) # 初始化推理引擎(如果模型加载成功) inference = None if model and tokenizer: try: from model.inference import ShimokumoInference inference = ShimokumoInference(config, model, tokenizer) except Exception as e: logger.error(f"推理引擎初始化失败: {e}") # 初始化功能模块 modules = init_modules(config, inference) # 计算启动耗时 init_duration = time.time() - start_time logger.info(f"初始化完成,耗时 {format_duration(init_duration)}") # 信号处理(优雅退出) def signal_handler(sig, frame): logger.info("收到退出信号,正在关闭...") print("\n\n霜云正在优雅退出...再见~") sys.exit(0) signal.signal(signal.SIGINT, signal_handler) signal.signal(signal.SIGTERM, signal_handler) # 选择运行模式 if args.interactive: # 交互式命令行模式 interactive_mode(config, inference, modules) else: # 服务器模式 try: from api.server import ShimokumoServer server = ShimokumoServer( config=config, inference_engine=inference, web_search_module=modules.get("web_search"), web_browser_module=modules.get("web_browser"), media_player_module=modules.get("media_player"), novel_creator_module=modules.get("novel_creator"), video_generator_module=modules.get("video_generator"), character_chat_module=modules.get("character_chat"), ) server.run() except ImportError as e: logger.error(f"缺少依赖: {e}") logger.error("请安装Flask: pip install flask") print("错误: 需要安装Flask来运行服务器模式。") print("请运行: pip install flask") print("或使用 --interactive 参数进入交互模式。") sys.exit(1) if __name__ == "__main__": main()