diff --git a/config.yml b/config.yml index 05ea060..d21bd60 100644 --- a/config.yml +++ b/config.yml @@ -38,7 +38,7 @@ react: # 协程池配置(各并行点并发度,缺失或 <1 时回退代码内默认值) pool: - kg_extract: 4 # 知识图谱:逐 chunk LLM 抽取 + kg_extract: 8 # 知识图谱:逐 chunk LLM 抽取(增加并发数加快构建速度) annotation_clause: 4 # 合同标注:逐条款 (recall+judge) annotation_dataset: 8 # 条款内:逐 dataset 召回 chat: 4 # 问答:retrieve 与 GraphEnhance 并行 diff --git a/kb/consts/consts.go b/kb/consts/consts.go index f2fd307..255a5d2 100644 --- a/kb/consts/consts.go +++ b/kb/consts/consts.go @@ -33,8 +33,8 @@ const ( ParsePollIntervalSeconds = 15 // 解析/标注任务轮询间隔 - DefaultLlmHttpTimeout = 600 // LLM HTTP 客户端超时(秒):本地模型生成慢,且多请求排队时响应可超分钟级 - MaxExtractTokens = 24576 // 知识抽取单次生成预算(思考链+JSON;24576≈实测自然输出 20 倍,4 并发下不超 131072 共享上下文) + DefaultLlmHttpTimeout = 600 // LLM HTTP 客户端超时(秒):本地模型生成慢,且多请求排队时响应可超分钟级 + MaxExtractTokens = 8192 // 知识抽取单次生成预算(JSON 输出实际 <500 token,8192 足够;减少输出 token 数可加快生成速度) EmbedBatchSize = 16 // 单次向量化请求的文本批量 @@ -57,7 +57,7 @@ const ( AnnoCandidateMaxChars = 2000 // 风险判定候选条文内容上限(需覆盖整 chunk,chunk 由多条短条文合并而成,截断会切掉块尾条文) // 协程池默认大小(config.yml pool 段缺失时兜底) - KgExtractPoolSize = 4 // 知识图谱:逐 chunk LLM 抽取 + KgExtractPoolSize = 8 // 知识图谱:逐 chunk LLM 抽取(增加并发数加快构建速度) AnnotationClausePoolSize = 4 // 合同标注:逐条款 (recall+judge) AnnotationDatasetPoolSize = 8 // 条款内:逐 dataset 召回 ChatPoolSize = 4 // 问答:retrieve 与 GraphEnhance 并行 diff --git a/scripts/start_models.sh b/scripts/start_models.sh index f3b4f68..387d303 100755 --- a/scripts/start_models.sh +++ b/scripts/start_models.sh @@ -2,6 +2,7 @@ # 启动本地 MLX 模型服务:chat + embedding 均由 oMLX(18080) 提供 # chat=Qwen3.5-9B-MLX-4bit(4-bit),embedding=BGE-M3(oMLX 原生 BERT 系支持) # 首次运行自动检测 .venv:缺失/损坏(如来自其他机器)时按 requirements.txt 重建,就绪则跳过安装 +# 根据设备内存大小自动调整上下文窗口和最大 token 数限制 # 用法: ./scripts/start_models.sh set -eu cd "$(dirname "$0")/.." @@ -9,8 +10,72 @@ ROOT=$(pwd) VENV="$ROOT/.venv" LOG_DIR="$HOME/.omlx/logs" CACHE_DIR="$HOME/.omlx/cache" +SETTINGS_FILE="$HOME/.omlx/settings.json" mkdir -p "$LOG_DIR" "$CACHE_DIR" +# 根据设备内存大小计算上下文窗口和最大 token 数 +calculate_memory_limits() { + # 获取系统总内存(字节) + local total_memory + total_memory=$(sysctl -n hw.memsize 2>/dev/null || echo 0) + + if [ "$total_memory" -eq 0 ]; then + echo "[warn] 无法检测系统内存,使用默认值" >&2 + echo "32768 32768" + return + fi + + local total_gb=$((total_memory / 1024 / 1024 / 1024)) + echo "[info] 检测到系统内存: ${total_gb}GB" >&2 + + # 计算可用内存(预留 10% 作为安全缓冲) + local available_memory=$((total_memory * 90 / 100)) + local available_gb=$((available_memory / 1024 / 1024 / 1024)) + + # 根据内存大小计算上下文窗口和最大 token 数 + # 考虑因素: + # 1. 模型加载:Qwen3.5-9B-4bit 约 5GB + # 2. KV 缓存:每 token 约 0.5MB(9B 模型) + # 3. 系统保留:预留 10% 内存 + # 4. 业务场景:合同审计/证物补全 + # - 知识库+知识图谱预检索过滤后才传递给模型 + # - 传递给模型的数据量本身不大(ToolResultMaxChars=400,HybridTopK=5) + # - 单次查询模式,不需要多轮对话缓存 + # - 16GB 机器 prefill 预算约 1900 token(2.9GB) + + local context_window=4096 + local max_tokens=2048 + local kv_cache_size="auto" + + if [ "$total_gb" -ge 32 ]; then + # 32GB+ 内存:适中配置即可 + context_window=4096 + max_tokens=2048 + kv_cache_size="4GB" + echo "[info] 内存充足,设置上下文窗口=${context_window},最大 token 数=${max_tokens},KV 缓存=${kv_cache_size}" >&2 + elif [ "$total_gb" -ge 24 ]; then + # 24GB 内存:标准配置 + context_window=4096 + max_tokens=2048 + kv_cache_size="3GB" + echo "[info] 中等内存配置,设置上下文窗口=${context_window},最大 token 数=${max_tokens},KV 缓存=${kv_cache_size}" >&2 + elif [ "$total_gb" -ge 16 ]; then + # 16GB 内存:优化配置(符合 prefill 预算) + context_window=4096 + max_tokens=2048 + kv_cache_size="2GB" + echo "[info] 基础内存配置,设置上下文窗口=${context_window},最大 token 数=${max_tokens},KV 缓存=${kv_cache_size}" >&2 + else + # 小于 16GB 内存:最小配置 + context_window=2048 + max_tokens=1024 + kv_cache_size="1GB" + echo "[info] 内存较小,设置上下文窗口=${context_window},最大 token 数=${max_tokens},KV 缓存=${kv_cache_size}" >&2 + fi + + echo "$context_window $max_tokens $kv_cache_size" +} + # .venv 不可跨机复制(pyvenv.cfg / bin 脚本 shebang 写死本机路径),须在本机重建 if [ ! -x "$VENV/bin/python" ]; then if [ -d "$VENV" ]; then @@ -35,6 +100,53 @@ else echo "[ok] .venv 已就绪,跳过安装" fi +# 根据硬件信息调整上下文窗口和最大 token 数限制 +if [ -f "$SETTINGS_FILE" ]; then + echo "[info] 检测硬件信息并调整内存限制..." + read -r context_window max_tokens kv_cache_size <<< "$(calculate_memory_limits)" + + # 更新 settings.json 中的 sampling 参数 + if command -v python3 >/dev/null 2>&1; then + python3 << EOF +import json +import sys + +try: + with open("$SETTINGS_FILE", 'r') as f: + settings = json.load(f) + + # 更新 sampling 参数 + if 'sampling' not in settings: + settings['sampling'] = {} + + old_context = settings['sampling'].get('max_context_window', 32768) + old_max_tokens = settings['sampling'].get('max_tokens', 32768) + + settings['sampling']['max_context_window'] = $context_window + settings['sampling']['max_tokens'] = $max_tokens + + # 更新 KV 缓存大小(如果配置了 cache 部分) + if 'cache' not in settings: + settings['cache'] = {} + + # 只有当 kv_cache_size 不是 "auto" 时才设置 + if "$kv_cache_size" != "auto": + settings['cache']['ssd_cache_max_size'] = "$kv_cache_size" + + with open("$SETTINGS_FILE", 'w') as f: + json.dump(settings, f, indent=2) + + print(f"[info] 已更新内存限制: 上下文窗口 {old_context} -> {$context_window}, 最大 token 数 {old_max_tokens} -> {$max_tokens}, KV 缓存: $kv_cache_size") +except Exception as e: + print(f"[warn] 更新 settings.json 失败: {e}", file=sys.stderr) +EOF + else + echo "[warn] 未找到 python3,跳过自动配置" >&2 + fi +else + echo "[info] settings.json 不存在,将使用默认配置" +fi + # oMLX 模型目录:chat(Qwen3.5-9B-MLX-4bit)+ embedding(bge-m3)均为实体目录,oMLX 按目录扫描注册模型 OMLX_MODEL_DIR="$ROOT/models/mlx/omlx"