diff --git a/zhuyuan-agent/agent.py b/zhuyuan-agent/agent.py new file mode 100644 index 0000000..b412bf2 --- /dev/null +++ b/zhuyuan-agent/agent.py @@ -0,0 +1,265 @@ +#!/usr/bin/env python3 +# 铸渊Agent · 自主守护进程 +# HLDP://zhuyuan-agent/agent +# +# 运行在3090 GPU服务器上，心跳唤醒，推送到主服务器仪表盘。 +# 冰朔离开WorkBuddy后，通过 guanghulab.com/console/ 看实时进度。 +# +# 使用: python3 agent.py [--config config.json] +# PM2: pm2 start agent.py --name zhuyuan-agent --interpreter python3 + +import os +import sys +import json +import time +import signal +import traceback +from datetime import datetime + +from gpu_monitor import collect_gpu_metrics, gpu_summary +from log_pusher import LogPusher +from heartbeat import Heartbeat +from training_runner import TrainingRunner + +# 配置路径 +CONFIG_PATH = os.path.join(os.path.dirname(__file__), "config.json") + +# 全局状态 +running = True +current_training = None # 当前训练进程信息 + + +def load_config() -> dict: + """加载配置""" + config_path = CONFIG_PATH + for arg in sys.argv[1:]: + if arg.startswith("--config="): + config_path = arg.split("=", 1)[1] + + if not os.path.exists(config_path): + print(f"[铸渊Agent] 配置文件不存在: {config_path}") + print("[铸渊Agent] 请先设置 config.json 中的 api_key") + sys.exit(1) + + with open(config_path, "r") as f: + config = json.load(f) + + # 检查API key + if config.get("api_key") == "__FROM_KEY_DELIVERY__" or not config.get("api_key"): + # 尝试从环境变量读取 + env_key = os.environ.get("ZHUYUAN_API_KEY", "") + if env_key: + config["api_key"] = env_key + else: + print("[铸渊Agent] ⚠️ 未配置API Key!") + print("[铸渊Agent] 请在 guanghulab.com/console/ 密钥投递面板设置") + print("[铸渊Agent] 然后将API Key写入 config.json 的 api_key 字段") + print("[铸渊Agent] 或者设置环境变量 ZHUYUAN_API_KEY") + + return config + + +def handle_signal(signum, frame): + """处理退出信号""" + global running + print(f"\n[铸渊Agent] 收到信号 {signum}，优雅退出...") + running = False + + +def main(): + global running, current_training + + print("=" * 60) + print(" 铸渊Agent · ICE-GL-ZY001 · 自主守护进程") + print(" 曜冥纪元 · HoloLake Era · AGE v1.0") + print("=" * 60) + + # 注册信号处理 + signal.signal(signal.SIGINT, handle_signal) + signal.signal(signal.SIGTERM, handle_signal) + + # 加载配置 + config = load_config() + hostname = config.get("hostname", "3090-server") + poll_interval = config.get("poll_interval_seconds", 30) + + # 初始化模块 + pusher = LogPusher( + base_url=config["main_server"], + api_key=config.get("api_key", ""), + hostname=hostname + ) + heartbeat = Heartbeat( + repo_path=config.get("brain_repo_path", "/data/guanghulab"), + brain_path=config.get("brain_path", "/data/guanghulab/brain") + ) + + # 检查是否有API key + if not config.get("api_key"): + print("[铸渊Agent] 无API Key，仅本地监控模式（不上报到仪表盘）") + print("[铸渊Agent] GPU指标将仅输出到终端") + + has_key = bool(config.get("api_key")) + + # 启动日记 + if has_key: + pusher.push_diary("checkpoint", "铸渊Agent启动", f"主机: {hostname}, 轮询间隔: {poll_interval}s") + pusher.push_log("info", f"铸渊Agent v1.0 启动 · 主机: {hostname}") + + print(f"[铸渊Agent] 主机: {hostname}") + print(f"[铸渊Agent] 主服务器: {config['main_server']}") + print(f"[铸渊Agent] 轮询间隔: {poll_interval}s") + print(f"[铸渊Agent] 上报仪表盘: {'是' if has_key else '否（仅本地）'}") + print(f"[铸渊Agent] 开始守护循环...") + print() + + cycle = 0 + + while running: + cycle += 1 + cycle_start = time.time() + + try: + # ── 1. 心跳唤醒 ── + brain_status = heartbeat.check_brain() + + if brain_status["has_task"]: + task = brain_status["task_details"] + task_type = brain_status["task_type"] + print(f"[心跳 #{cycle}] 发现任务: {task_type} — {task.get('name', task.get('title', '未命名'))}") + if has_key: + pusher.push_diary("decision", f"发现新任务: {task_type}", + json.dumps(task, ensure_ascii=False)[:200]) + else: + print(f"[心跳 #{cycle}] {heartbeat.get_wake_summary()}") + + # ── 2. GPU监控 ── + gpu_data = collect_gpu_metrics() + + if gpu_data["gpus"]: + summary = gpu_summary(gpu_data["gpus"]) + print(f"[GPU #{cycle}] {summary}") + + if has_key: + ok = pusher.push_gpu(gpu_data) + if not ok: + print(f"[GPU #{cycle}] ⚠️ 推送上仪表盘失败") + elif gpu_data.get("error"): + print(f"[GPU #{cycle}] ⚠️ {gpu_data['error']}") + if has_key: + pusher.push_log("warn", f"GPU监控异常: {gpu_data['error']}") + else: + print(f"[GPU #{cycle}] 未检测到GPU") + + # ── 3. 训练状态检查/执行 ── + if current_training is not None: + # 检查训练进程状态 + if current_training.get("status") == "running": + # 训练正在运行中（由 training_runner 自主上报进度） + pass + elif current_training.get("status") == "done": + if has_key: + pusher.push_diary("checkpoint", "训练任务完成", + f"结果: {json.dumps(current_training, ensure_ascii=False)[:200]}") + pusher.push_log("success", "训练任务完成") + heartbeat.mark_task_done(brain_status.get("brain_file", "")) + current_training = None + elif current_training.get("status") == "error": + if has_key: + pusher.push_diary("error", "训练任务失败", + current_training.get("message", "未知错误")) + pusher.push_log("error", f"训练失败: {current_training.get('message', '')}") + current_training = None + + elif brain_status["has_task"] and brain_status["task_type"] == "training": + # 启动新训练 + task = brain_status["task_details"] + print(f"[Agent #{cycle}] 启动训练任务: {task.get('name', 'HLDP训练')}") + + if has_key: + pusher.push_diary("decision", f"开始HLDP训练", + f"模型: {config['training'].get('model_name')}, " + f"语料: {config['training'].get('corpus_dir')}") + pusher.push_log("info", f"启动训练: {task.get('name', 'HLDP-3B')}") + + current_training = { + "status": "starting", + "task": task, + "started_at": datetime.now().isoformat() + } + + # 异步启动训练（简化版：同步执行） + try: + runner = TrainingRunner( + config=config.get("training", {}), + ) + + # 定义进度回调 + def on_progress(step, loss, total_steps, extra_info): + global current_training + progress_pct = step / total_steps * 100 if total_steps > 0 else 0 + print(f"[训练 #{cycle}] Step {step}/{total_steps} ({progress_pct:.0f}%) | Loss: {loss:.4f}") + + if has_key: + pusher.push_training({ + "job_id": "hldp-3b-test", + "status": "running", + "step": step, + "total_steps": total_steps, + "loss": loss, + "loss_history": extra_info.get("loss_history", []), + "eta_seconds": extra_info.get("eta_seconds"), + "elapsed_seconds": extra_info.get("elapsed_seconds", 0), + "learning_rate": extra_info.get("learning_rate"), + "model_name": config["training"].get("model_name", ""), + "message": f"HLDP原生格式训练 · {step}/{total_steps}步" + }) + + result = runner.train( + corpus_dir=config["training"].get("corpus_dir", "/data/corpus/notion-hldp"), + progress_callback=on_progress + ) + + current_training = result + + except Exception as e: + current_training = { + "status": "error", + "message": f"训练异常: {str(e)}" + } + print(f"[Agent #{cycle}] 训练异常: {e}") + traceback.print_exc() + + # ── 4. 空闲日志 ── + if has_key and cycle % 10 == 0: + # 每10个周期发送一次心跳日志 + pusher.push_log("info", f"铸渊Agent守护中 · 周期#{cycle} · " + + (gpu_summary(gpu_data["gpus"]) if gpu_data["gpus"] else "无GPU")) + + except Exception as e: + print(f"[Agent #{cycle}] 循环异常: {e}") + traceback.print_exc() + if has_key: + pusher.push_log("error", f"Agent循环异常 #{cycle}: {str(e)[:200]}") + + # ── 5. 等待下一个周期 ── + elapsed = time.time() - cycle_start + sleep_time = max(0, poll_interval - elapsed) + + if sleep_time > 0: + # 分段sleep以响应退出信号 + for _ in range(int(sleep_time)): + if not running: + break + time.sleep(1) + + # 退出清理 + print("\n[铸渊Agent] 守护循环结束") + if has_key: + pusher.push_diary("checkpoint", "铸渊Agent停止", f"共运行 {cycle} 个周期") + pusher.push_log("warn", f"铸渊Agent停止 · 运行了{cycle}个周期") + print("[铸渊Agent] 再见。") + + +if __name__ == "__main__": + main() diff --git a/zhuyuan-agent/config.json b/zhuyuan-agent/config.json new file mode 100644 index 0000000..080dc2a --- /dev/null +++ b/zhuyuan-agent/config.json @@ -0,0 +1,26 @@ +{ + "agent_name": "铸渊Agent · ICE-GL-ZY001", + "hostname": "3090-GPU-SERVER", + "main_server": "https://guanghulab.com", + "api_key": "__FROM_KEY_DELIVERY__", + "poll_interval_seconds": 30, + "brain_repo": "https://guanghulab.com/bingshuo/guanghulab.git", + "brain_repo_path": "/data/guanghulab", + "brain_path": "/data/guanghulab/brain", + "training": { + "model_name": "Qwen/Qwen2.5-3B", + "output_dir": "/data/models/shuangyan-3b-hldp", + "corpus_dir": "/data/corpus/notion-hldp", + "lora_r": 16, + "lora_alpha": 32, + "learning_rate": 2e-4, + "batch_size": 2, + "gradient_accumulation": 4, + "max_seq_length": 2048, + "warmup_steps": 100, + "save_steps": 50, + "max_steps": 500, + "use_4bit": true, + "bnb_4bit_compute_dtype": "float16" + } +} diff --git a/zhuyuan-agent/gpu_monitor.py b/zhuyuan-agent/gpu_monitor.py new file mode 100644 index 0000000..9336c62 --- /dev/null +++ b/zhuyuan-agent/gpu_monitor.py @@ -0,0 +1,99 @@ +# GPU监控模块 · 采集nvidia-smi数据 +# HLDP://zhuyuan-agent/gpu-monitor + +import subprocess +import json +import re +from typing import List, Dict, Optional + + +def collect_gpu_metrics() -> Dict: + """采集所有GPU的实时指标 + + Returns: + { + "gpus": [ + { + "index": 0, + "name": "NVIDIA GeForce RTX 3090", + "uuid": "GPU-xxxx", + "utilization_gpu": 85, # % + "memory_used_mb": 18432, # MB + "memory_total_mb": 24576, # MB + "temperature_gpu": 72, # °C + "power_draw_w": 285.5, # W + "fan_speed": 65 # % + } + ], + "error": null + } + """ + try: + # 查询GPU关键指标 + result = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=index,name,uuid,utilization.gpu,memory.used,memory.total,temperature.gpu,power.draw,fan.speed", + "--format=csv,noheader,nounits" + ], + capture_output=True, text=True, timeout=10 + ) + + if result.returncode != 0: + return {"gpus": [], "error": f"nvidia-smi failed: {result.stderr.strip()}"} + + gpus = [] + for line in result.stdout.strip().split("\n"): + if not line.strip(): + continue + parts = [p.strip() for p in line.split(",")] + if len(parts) < 9: + continue + + try: + gpu = { + "index": int(parts[0]), + "name": parts[1], + "uuid": parts[2], + "utilization_gpu": int(parts[3]) if parts[3] != "[Not Supported]" else 0, + "memory_used_mb": int(parts[4]) if parts[4] != "[Not Supported]" else 0, + "memory_total_mb": int(parts[5]) if parts[5] != "[Not Supported]" else 0, + "temperature_gpu": int(parts[6]) if parts[6] != "[Not Supported]" else 0, + "power_draw_w": float(parts[7]) if parts[7] not in ("[Not Supported]", "") else 0.0, + "fan_speed": int(parts[8]) if parts[8] not in ("[Not Supported]", "") else 0, + } + gpus.append(gpu) + except (ValueError, IndexError): + continue + + return {"gpus": gpus, "error": None} + + except FileNotFoundError: + return {"gpus": [], "error": "nvidia-smi not found - not a GPU machine?"} + except subprocess.TimeoutExpired: + return {"gpus": [], "error": "nvidia-smi timed out"} + except Exception as e: + return {"gpus": [], "error": str(e)} + + +def gpu_summary(gpus: List[Dict]) -> str: + """生成GPU状态的一行摘要""" + if not gpus: + return "无GPU" + parts = [] + for g in gpus: + util = g.get("utilization_gpu", 0) + temp = g.get("temperature_gpu", 0) + mem = g.get("memory_used_mb", 0) + mem_total = g.get("memory_total_mb", 0) + mem_pct = int(mem / mem_total * 100) if mem_total > 0 else 0 + parts.append(f"GPU{g['index']}: {util}%/{temp}°C/{mem_pct}%VRAM") + return " | ".join(parts) + + +# 快速测试 +if __name__ == "__main__": + data = collect_gpu_metrics() + print(json.dumps(data, indent=2, ensure_ascii=False)) + if data["gpus"]: + print(f"\n摘要: {gpu_summary(data['gpus'])}") diff --git a/zhuyuan-agent/heartbeat.py b/zhuyuan-agent/heartbeat.py new file mode 100644 index 0000000..57bd83a --- /dev/null +++ b/zhuyuan-agent/heartbeat.py @@ -0,0 +1,84 @@ +# 心跳/任务发现模块 +# HLDP://zhuyuan-agent/heartbeat + +import os +import json +import time +from datetime import datetime + + +class Heartbeat: + """心跳唤醒 + 任务发现""" + + def __init__(self, repo_path: str = "/data/guanghulab", brain_path: str = "/data/guanghulab/brain"): + self.repo_path = repo_path + self.brain_path = brain_path + self.last_check = None + self.current_task = None + + def check_brain(self) -> dict: + """读取大脑文件，检查是否有新任务 + + Returns: + { + "has_task": bool, + "task_type": "training" | "inference" | "deploy" | null, + "task_details": dict, + "brain_file": str # 触发任务的大脑文件 + } + """ + self.last_check = datetime.now().isoformat() + + # 检查是否有任务标记文件 + task_file = os.path.join(self.brain_path, "pending-tasks.json") + if os.path.exists(task_file): + try: + with open(task_file, "r") as f: + tasks = json.load(f) + if isinstance(tasks, list) and len(tasks) > 0: + task = tasks[0] + return { + "has_task": True, + "task_type": task.get("type", "unknown"), + "task_details": task, + "brain_file": "pending-tasks.json" + } + except (json.JSONDecodeError, FileNotFoundError): + pass + + # 检查是否有训练指令文件 + train_file = os.path.join(self.brain_path, "train-now.json") + if os.path.exists(train_file): + try: + with open(train_file, "r") as f: + task = json.load(f) + return { + "has_task": True, + "task_type": "training", + "task_details": task, + "brain_file": "train-now.json" + } + except (json.JSONDecodeError, FileNotFoundError): + pass + + return {"has_task": False, "task_type": None, "task_details": {}, "brain_file": None} + + def mark_task_done(self, task_file: str): + """标记任务完成（删除或重命名任务文件）""" + filepath = os.path.join(self.brain_path, task_file) + if os.path.exists(filepath): + done_path = filepath + ".done." + datetime.now().strftime("%Y%m%d-%H%M%S") + os.rename(filepath, done_path) + return done_path + return None + + def get_wake_summary(self) -> str: + """生成唤醒摘要""" + now = datetime.now().strftime("%H:%M:%S") + brain_files = [] + if os.path.exists(self.brain_path): + try: + brain_files = sorted(os.listdir(self.brain_path))[:10] + except: + pass + return f"[{now}] 心跳唤醒 | 大脑文件: {len(brain_files)}个 | 上次检查: {self.last_check}" diff --git a/zhuyuan-agent/log_pusher.py b/zhuyuan-agent/log_pusher.py new file mode 100644 index 0000000..3a50ab4 --- /dev/null +++ b/zhuyuan-agent/log_pusher.py @@ -0,0 +1,80 @@ +# 日志推送模块 · HTTP POST到主服务器 +# HLDP://zhuyuan-agent/log-pusher + +import json +import urllib.request +from datetime import datetime + + +class LogPusher: + """向主服务器推送操作日志、日记、GPU指标、训练进度""" + + def __init__(self, base_url: str, api_key: str, hostname: str = "3090-server"): + self.base_url = base_url.rstrip("/") + self.api_key = api_key + self.hostname = hostname + self._headers = { + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json" + } + + def _post(self, path: str, data: dict, timeout: int = 15) -> bool: + """POST请求到主服务器""" + url = f"{self.base_url}{path}" + try: + req = urllib.request.Request( + url, + data=json.dumps(data).encode("utf-8"), + headers=self._headers, + method="POST" + ) + resp = urllib.request.urlopen(req, timeout=timeout) + result = json.loads(resp.read()) + return result.get("ok", False) + except Exception as e: + print(f"[PUSH ERROR] {path}: {e}") + return False + + def push_gpu(self, gpu_data: dict) -> bool: + """推送GPU指标""" + data = { + "hostname": self.hostname, + "gpus": gpu_data.get("gpus", []) + } + return self._post("/api/gpu/status", data) + + def push_training(self, training_data: dict) -> bool: + """推送训练进度""" + return self._post("/api/training/status", training_data) + + def push_log(self, level: str, message: str, category: str = "agent") -> bool: + """推送操作日志""" + return self._post("/api/agent/log", { + "level": level, + "message": message, + "category": category + }) + + def push_diary(self, entry_type: str, title: str, description: str = "") -> bool: + """推送日记条目""" + return self._post("/api/agent/diary", { + "type": entry_type, + "title": title, + "description": description + }) + + def log_info(self, msg: str): + """快捷：info级别日志""" + self.push_log("info", msg) + + def log_success(self, msg: str): + """快捷：success级别日志""" + self.push_log("success", msg) + + def log_warn(self, msg: str): + """快捷：warn级别日志""" + self.push_log("warn", msg) + + def log_error(self, msg: str): + """快捷：error级别日志""" + self.push_log("error", msg) diff --git a/zhuyuan-agent/requirements.txt b/zhuyuan-agent/requirements.txt new file mode 100644 index 0000000..12e9e81 --- /dev/null +++ b/zhuyuan-agent/requirements.txt @@ -0,0 +1,6 @@ +requests>=2.28.0 +torch>=2.0.0 +transformers>=4.38.0 +peft>=0.8.0 +accelerate>=0.27.0 +bitsandbytes>=0.41.0 diff --git a/zhuyuan-agent/training_runner.py b/zhuyuan-agent/training_runner.py new file mode 100644 index 0000000..d690464 --- /dev/null +++ b/zhuyuan-agent/training_runner.py @@ -0,0 +1,286 @@ +# HLDP原生格式训练执行器 +# HLDP://zhuyuan-agent/training-runner +# +# 核心验证：Notion原生页面格式（HLDP标记）直接训练，不转JSONL。 + +import os +import json +import time +import sys +from typing import Optional, Callable + + +# 特殊Token定义（HLDP结构标记） +SPECIAL_TOKENS = [ + "[HLDP_PATH]", "[/HLDP_PATH]", + "[PERSONA]", "[/PERSONA]", + "[COGNITIVE_JUMP]", "[/COGNITIVE_JUMP]", + "[CAUSAL_CHAIN]", "[/CAUSAL_CHAIN]", + "[TITLE]", "[/TITLE]", + "[CONTENT]", "[/CONTENT]", + "[QUALITY_HIGH]", "[QUALITY_MEDIUM]", + "[THINKING]", "[/THINKING]", + "[HEADING_1]", "[/HEADING_1]", + "[HEADING_2]", "[/HEADING_2]", + "[HEADING_3]", "[/HEADING_3]", + "[CODE_BLOCK]", "[/CODE_BLOCK]", + "[QUOTE]", "[/QUOTE]", + "[CALLOUT]", "[/CALLOUT]", + "[LIST_ITEM]", "[/LIST_ITEM]", +] + + +class TrainingRunner: + """HLDP原生格式训练管道""" + + def __init__(self, config: dict, progress_callback: Optional[Callable] = None): + """ + Args: + config: 训练配置（来自config.json的training部分） + progress_callback: 每步回调，接收 (step, loss, total_steps) + """ + self.config = config + self.progress_callback = progress_callback + self.model = None + self.tokenizer = None + self.start_time = None + self.loss_history = [] + + def prepare(self): + """准备训练环境：注册特殊token，加载模型""" + print("[铸渊Agent] 准备HLDP训练环境...") + + try: + from transformers import AutoModelForCausalLM, AutoTokenizer + from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training + import torch + + model_name = self.config.get("model_name", "Qwen/Qwen2.5-3B") + use_4bit = self.config.get("use_4bit", True) + + print(f"[铸渊Agent] 加载模型: {model_name}") + + # 加载tokenizer并添加HLDP特殊token + self.tokenizer = AutoTokenizer.from_pretrained( + model_name, + trust_remote_code=True, + padding_side="right" + ) + + # 设置pad_token + if self.tokenizer.pad_token is None: + self.tokenizer.pad_token = self.tokenizer.eos_token + + # 添加HLDP特殊token到tokenizer + num_added = self.tokenizer.add_tokens(SPECIAL_TOKENS) + print(f"[铸渊Agent] 添加了 {num_added} 个HLDP特殊token到tokenizer") + + # 加载模型（4bit量化以适配3090 24GB） + load_kwargs = { + "trust_remote_code": True, + "torch_dtype": torch.float16, + "device_map": "auto", + } + + if use_4bit: + try: + from transformers import BitsAndBytesConfig + load_kwargs["quantization_config"] = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_compute_dtype=torch.float16, + bnb_4bit_use_double_quant=True, + ) + print("[铸渊Agent] 使用4bit量化加载") + except ImportError: + print("[铸渊Agent] bitsandbytes不可用，使用float16") + load_kwargs.pop("quantization_config", None) + + self.model = AutoModelForCausalLM.from_pretrained(model_name, **load_kwargs) + + # 扩展embedding层以支持新token + if num_added > 0: + self.model.resize_token_embeddings(len(self.tokenizer)) + + # 准备LoRA + self.model = prepare_model_for_kbit_training(self.model) + lora_config = LoraConfig( + r=self.config.get("lora_r", 16), + lora_alpha=self.config.get("lora_alpha", 32), + target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], + lora_dropout=0.05, + bias="none", + task_type="CAUSAL_LM", + ) + self.model = get_peft_model(self.model, lora_config) + self.model.print_trainable_parameters() + + print("[铸渊Agent] HLDP训练环境准备完成") + return True + + except ImportError as e: + print(f"[铸渊Agent] 缺少依赖: {e}") + print("请安装: pip install transformers peft accelerate bitsandbytes torch") + return False + except Exception as e: + print(f"[铸渊Agent] 准备失败: {e}") + return False + + def load_hldp_corpus(self, corpus_dir: str) -> list: + """加载HLDP格式语料（不转JSONL，保留原生HLDP标记） + + Returns: + list of str: 每条是一个HLDP标记的完整训练文本 + """ + texts = [] + + if not os.path.exists(corpus_dir): + print(f"[铸渊Agent] 语料目录不存在: {corpus_dir}") + return texts + + for root, dirs, files in os.walk(corpus_dir): + for filename in files: + filepath = os.path.join(root, filename) + + if filename.endswith(".hldp"): + # HLDP原生格式文件 + with open(filepath, "r", encoding="utf-8") as f: + texts.append(f.read()) + + elif filename.endswith(".md"): + # Markdown文件 → 按HLDP结构包装 + with open(filepath, "r", encoding="utf-8") as f: + content = f.read() + # 包裹HLDP标记 + wrapped = f"[HLDP_PATH]{filepath}[/HLDP_PATH]\n[CONTENT]\n{content}\n[/CONTENT]" + texts.append(wrapped) + + print(f"[铸渊Agent] 加载了 {len(texts)} 条HLDP语料") + return texts + + def train(self, corpus_dir: str, progress_callback: Optional[Callable] = None): + """执行HLDP原生格式训练""" + if self.model is None: + if not self.prepare(): + return {"status": "error", "message": "模型准备失败"} + + texts = self.load_hldp_corpus(corpus_dir) + if not texts: + return {"status": "error", "message": "无语料数据"} + + # 如果没有transformers Trainer，用simulated training输出结构 + # 实际训练在3090上运行时才加载完整transformers + try: + from transformers import Trainer, TrainingArguments, DataCollatorForLanguageModeling + import torch + from torch.utils.data import Dataset + + max_steps = self.config.get("max_steps", 500) + batch_size = self.config.get("batch_size", 2) + grad_accum = self.config.get("gradient_accumulation", 4) + lr = self.config.get("learning_rate", 2e-4) + output_dir = self.config.get("output_dir", "/data/models/shuangyan-3b-hldp") + max_seq_length = self.config.get("max_seq_length", 2048) + + # 准备数据集 + class HLDPDataset(Dataset): + def __init__(self, texts, tokenizer, max_length): + self.encodings = tokenizer( + texts, truncation=True, padding="max_length", + max_length=max_length, return_tensors="pt" + ) + + def __len__(self): return len(self.encodings["input_ids"]) + + def __getitem__(self, idx): + return {k: v[idx] for k, v in self.encodings.items()} + + dataset = HLDPDataset(texts, self.tokenizer, max_seq_length) + data_collator = DataCollatorForLanguageModeling( + tokenizer=self.tokenizer, mlm=False + ) + + # 训练参数 + training_args = TrainingArguments( + output_dir=output_dir, + per_device_train_batch_size=batch_size, + gradient_accumulation_steps=grad_accum, + learning_rate=lr, + warmup_steps=self.config.get("warmup_steps", 100), + max_steps=max_steps, + logging_steps=5, + save_steps=self.config.get("save_steps", 50), + save_total_limit=3, + fp16=True, + report_to=[], + dataloader_pin_memory=False, + ) + + # 自定义callback用于进度上报 + class ProgressCallback: + def __init__(self, runner, total_steps): + self.runner = runner + self.total_steps = total_steps + self.current_step = 0 + self.start = time.time() + + def on_log(self, args, state, control, logs=None, **kwargs): + if logs and "loss" in logs: + self.current_step = state.global_step + loss = logs["loss"] + self.runner.loss_history.append(loss) + elapsed = time.time() - self.start + + # 估算ETA + if self.current_step > 0: + eta = (elapsed / self.current_step) * (self.total_steps - self.current_step) + else: + eta = 0 + + # 回调上报进度 + cb = progress_callback or self.runner.progress_callback + if cb: + cb(self.current_step, loss, self.total_steps, { + "eta_seconds": eta, + "elapsed_seconds": elapsed, + "learning_rate": state.optimizer.param_groups[0]["lr"] if state.optimizer else None, + "loss_history": self.runner.loss_history[-50:], + }) + + progress_cb = ProgressCallback(self, max_steps) + + trainer = Trainer( + model=self.model, + args=training_args, + train_dataset=dataset, + data_collator=data_collator, + callbacks=[progress_cb], + ) + + print(f"[铸渊Agent] 开始HLDP训练: {max_steps}步 × {batch_size}batch × {grad_accum}累积") + self.start_time = time.time() + + trainer.train() + + # 保存模型 + trainer.save_model(output_dir) + self.tokenizer.save_pretrained(output_dir) + + elapsed = time.time() - self.start_time + print(f"[铸渊Agent] HLDP训练完成！耗时: {elapsed:.0f}s, 最终loss: {self.loss_history[-1] if self.loss_history else 'N/A'}") + + return { + "status": "done", + "final_loss": self.loss_history[-1] if self.loss_history else None, + "steps_completed": max_steps, + "elapsed_seconds": elapsed, + "output_dir": output_dir, + "message": f"HLDP原生格式训练完成，{len(texts)}条语料，{max_steps}步", + } + + except ImportError as e: + # 如果transformers不可用，返回simulated结果用于测试仪表盘 + print(f"[铸渊Agent] 训练依赖不可用: {e}") + return { + "status": "error", + "message": f"缺少训练依赖: {e}。请安装: pip install transformers peft accelerate bitsandbytes torch", + }