From 9f85944646f7ba2c711fc713f838f67647d8337a Mon Sep 17 00:00:00 2001 From: shuiyuPyLab Date: Sun, 9 Aug 2026 10:17:30 +0800 Subject: [PATCH] =?UTF-8?q?chore:=20=E9=A1=B9=E7=9B=AE=E5=B7=B2=E8=BF=81?= =?UTF-8?q?=E7=A7=BB=E8=87=B3=20Gitee/GitHub=EF=BC=8C=E4=BB=85=E4=BF=9D?= =?UTF-8?q?=E7=95=99=E8=BF=81=E7=A7=BB=E8=AF=B4=E6=98=8E=20(tombstone)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .cnb.yml | 21 - .gitattributes | 2 - .gitignore | 84 -- LICENSE | 21 - README.md | 3 +- api_correction.py | 172 ---- build_vocab.py | 134 --- chat.html | 284 ------ chat_server.py | 143 --- distill_daemon.py | 197 ---- gen_pretrain.py | 484 ---------- llm_chat_improved.py | 1330 --------------------------- scripts/api_correction.py | 148 --- scripts/daily_data_gen.py | 221 ----- scripts/generate_batch.py | 57 -- scripts/generate_diverse.py | 73 -- scripts/generate_diverse_single.py | 132 --- scripts/generate_fresh_qa.py | 240 ----- scripts/generate_massive_general.py | 389 -------- scripts/generate_multiturn.py | 346 ------- scripts/generate_real_multiturn.py | 337 ------- scripts/regenerate_all_high_temp.py | 126 --- scripts/train_cnb.py | 204 ---- setup_cloud.sh | 90 -- templates/chat.html | 309 ------- train.py | 495 ---------- train_cloud.py | 437 --------- train_max.py | 662 ------------- web_server.py | 128 --- xiaoli_next.py | 422 --------- 30 files changed, 1 insertion(+), 7690 deletions(-) delete mode 100644 .cnb.yml delete mode 100644 .gitattributes delete mode 100644 .gitignore delete mode 100644 LICENSE delete mode 100644 api_correction.py delete mode 100644 build_vocab.py delete mode 100644 chat.html delete mode 100644 chat_server.py delete mode 100644 distill_daemon.py delete mode 100644 gen_pretrain.py delete mode 100644 llm_chat_improved.py delete mode 100644 scripts/api_correction.py delete mode 100644 scripts/daily_data_gen.py delete mode 100644 scripts/generate_batch.py delete mode 100644 scripts/generate_diverse.py delete mode 100644 scripts/generate_diverse_single.py delete mode 100644 scripts/generate_fresh_qa.py delete mode 100644 scripts/generate_massive_general.py delete mode 100644 scripts/generate_multiturn.py delete mode 100644 scripts/generate_real_multiturn.py delete mode 100644 scripts/regenerate_all_high_temp.py delete mode 100644 scripts/train_cnb.py delete mode 100644 setup_cloud.sh delete mode 100644 templates/chat.html delete mode 100644 train.py delete mode 100644 train_cloud.py delete mode 100644 train_max.py delete mode 100644 web_server.py delete mode 100644 xiaoli_next.py diff --git a/.cnb.yml b/.cnb.yml deleted file mode 100644 index 9e9807d..0000000 --- a/.cnb.yml +++ /dev/null @@ -1,21 +0,0 @@ -$: - vscode: - - docker: - build: .ide/Dockerfile - runner: - cpus: 8 - tags: cnb:arch:amd64:gpu - services: - - vscode - - docker - env: - CUDA_VISIBLE_DEVICES: "0" - PYTHONUNBUFFERED: "1" - stages: - - name: 开始训练 - script: | - cd /workspace - nvidia-smi - echo "=== GPU就绪，开始训练 ===" - python3 scripts/train_cnb.py --epochs 100 --batch_size 16 - echo "=== 训练完成 ===" \ No newline at end of file diff --git a/.gitattributes b/.gitattributes deleted file mode 100644 index e700c3b..0000000 --- a/.gitattributes +++ /dev/null @@ -1,2 +0,0 @@ -*.json filter=lfs diff=lfs merge=lfs -text -llm_model_* filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore deleted file mode 100644 index 59a948d..0000000 --- a/.gitignore +++ /dev/null @@ -1,84 +0,0 @@ -# 已跟踪 -# 缓存和日志 -__pycache__/ -*.pyc -*.log -/tmp/ - -# 模型权重 — 只追踪发布版本(精简推理版,不含optimizer,约60MB) -*.pth -!/llm_model_v2_resume_epoch_7.pth -!/xiaoli-llm-chat/xiaoli_v2_best_inference.pth -!/xiaoli-llm-chat/xiaoli_v2_final_inference.pth - -# 中间输出目录（生成脚本的输出） -high_temp_outputs/ -diverse_answers/ -fresh_qa_outputs/ -multiturn_outputs/ -checkpoints_v2/ -old_ckpts/ -backup_old_arch/ - -# 旧版废弃数据 -xiaoli_*.json -mega_interim_*.json -essay_interim_*.json -distill_*.json -cot_*.json -memory_long*.json -multiturn_batch_*.json -real_multiturn*.json -mega_questions*.json -math_questions*.json -essay_questions*.json -cot_candidates*.json -model_test_results*.json -online_distill_report*.json - -# 测试/临时文件 -test_*.json -test_*.py -/var/folders - -# 日志 -batch_runner*.log -mega_*.log -mimo_*.log -knowledge_*.log -train_*.log -web_server*.log -# CNB 配置 — 强制跟踪关键文件 -.ide/ -!/.ide/Dockerfile -!/.cnb.yml - -# 训练输出目录 -general_outputs/ -high_temp_outputs/ -multiturn_outputs/ - -# 本地临时 -fun_test.py -*.log - -# 打包/发布产物 -xiaoli_package/ -*.zip - -# 上传/测试脚本 -upload_*.py -_gofile_* -_t.* -_up_test.* - -# 旧词表备份 -vocab_old.json -vocab_backup*.json - -# 数据目录（太大不进git） -data/ - -# 训练checkpoint（不进git，通过release发布） -xiaoli_cnb_best.pth -xiaoli_cnb_final.pth diff --git a/LICENSE b/LICENSE deleted file mode 100644 index 7203af6..0000000 --- a/LICENSE +++ /dev/null @@ -1,21 +0,0 @@ -MIT License - -Copyright (c) 2026 shuiyu - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. diff --git a/README.md b/README.md index 51b3c6a..3e6e785 100644 --- a/README.md +++ b/README.md @@ -2,8 +2,7 @@ 你好呀～这个仓库已经搬到更好用的地方啦。 -为了让代码管理更清爽、多平台同步更省心，我们把原先分散的镜像统一收拢到 -**Gitee** 与 **GitHub** 主仓维护，这里就只保留这张"迁移小纸条"咯。 +为了让代码管理更清爽、多平台同步更省心，我们把原先分散的镜像统一收拢到主仓维护，这里就只保留这张"迁移小纸条"咯。 📦 新家在这里： - Gitee：https://gitee.com/shuiyu1123/xiaoli-llm-chat diff --git a/api_correction.py b/api_correction.py deleted file mode 100644 index d230cfb..0000000 --- a/api_correction.py +++ /dev/null @@ -1,172 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -API生成模块 — Astron 5并发，稳定不超限 -""" -import json, requests, time, os, threading -from concurrent.futures import ThreadPoolExecutor, as_completed - -ASTRON_URL = "https://maas-coding-api.cn-huabei-1.xf-yun.com/v2/chat/completions" -ASTRON_KEY = os.environ.get("ASTRON_KEY", "") -ASTRON_MODEL = "astron-code-latest" - -SYSTEM_PROMPT = """你是一个名叫小狸的8岁中国小女孩。 - -你的人设： -- 8岁女孩子，有点害羞但很可爱 -- 说话简单自然，符合小朋友的语气 -- 回答问题要简短，1-2句话就够了 -- 不要用"喵"结尾，不要过度卖萌 - -回答要求： -1. 语气自然，像真实的小朋友说话 -2. 不懂就说"不知道" -3. 回答简洁但要说清楚，不要讲一半就停""" - - -def call_astron(q, temperature=0.85, system_prompt=None, max_tokens=8192): - """调用Astron API，支持自定义system prompt""" - if system_prompt is None: - system_prompt = SYSTEM_PROMPT - for attempt in range(3): - try: - headers = {"Content-Type": "application/json", "Authorization": f"Bearer {ASTRON_KEY}"} - payload = { - "model": ASTRON_MODEL, - "messages": [ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": q} - ], - "temperature": temperature, - "max_tokens": max_tokens, - } - resp = requests.post(ASTRON_URL, headers=headers, json=payload, timeout=120) - if resp.status_code == 200: - content = resp.json()['choices'][0]['message']['content'].strip() - if content: - return content - else: - time.sleep(1) - except Exception: - time.sleep(2) - return None - - -def batch_generate(questions, temperature=0.85, max_workers=5): - """Astron 5并发 — 明确限制，稳定可靠""" - total = len(questions) - results = [] - - def call(q): - ans = call_astron(q, temperature) - if ans: - return q, ans - return None - - with ThreadPoolExecutor(max_workers=max_workers) as executor: - futures = {executor.submit(call, item): i for i, item in enumerate(questions)} - done = 0 - t0 = time.time() - for f in as_completed(futures): - done += 1 - r = f.result() - if r: - results.append({"question": r[0], "answer": r[1]}) - if done % 20 == 0 or done == total: - elapsed = max(time.time() - t0, 1) - rate = len(results) / elapsed * 60 - print(f" Astron: {done}/{total} (已获{len(results)}条, {rate:.0f}条/分)", end="\r") - - elapsed = time.time() - t0 - rate = len(results) / max(elapsed, 1) * 60 - print(f" Astron: {total}/{total} → 共{len(results)}条 ({elapsed:.0f}s, {rate:.0f}条/分)") - return results - - -# 保持旧函数名兼容 -batch_generate_astron = batch_generate -batch_generate_mimo = batch_generate -batch_generate_direct = batch_generate -batch_generate_dual = batch_generate - - -def batch_convert(items, progress=None, max_workers=5): - """修复版 batch_convert — 将 (q, ans) 对用 API 优化后返回""" - if not items: - return [] - total = len(items) - results = [] - - def convert(item): - q, model_ans = item - prompt = f"问题：{q}\n模型回答：{model_ans}\n\n请用8岁小狸的语气，重新生成更好的回答。保持简短自然，1-2句话。" - ans = call_astron(prompt, temperature=0.3) - if ans and len(ans) >= 3: - return {"question": q, "answer": ans} - return None - - with ThreadPoolExecutor(max_workers=max_workers) as executor: - futures = {executor.submit(convert, item): i for i, item in enumerate(items)} - done = 0 - t0 = time.time() - for f in as_completed(futures): - done += 1 - r = f.result() - if r: - results.append(r) - if progress: - progress(done, total) - if done % 50 == 0 or done == total: - elapsed = max(time.time() - t0, 1) - rate = len(results) / elapsed * 60 - print(f" API转换: {done}/{total} (+{len(results)}条, {rate:.0f}条/分)", end="\r") - - elapsed = time.time() - t0 - rate = len(results) / max(elapsed, 1) * 60 - print(f" API转换: {total}/{total} → 共{len(results)}条 ({elapsed:.0f}s, {rate:.0f}条/分)") - return results - - -def batch_generate_multi_temp(questions, temperatures=None, max_workers=5): - """修复版 — 用多种温度生成多样化回答""" - if temperatures is None: - temperatures = [0.7, 0.85, 1.0] - if not questions: - return [] - - total = len(questions) * len(temperatures) - results = [] - - def call_with_temp(q, temp): - ans = call_astron(q, temperature=temp) - if ans: - return {"question": q, "answer": ans, "temperature": temp} - return None - - tasks = [(q, t) for q in questions for t in temperatures] - with ThreadPoolExecutor(max_workers=max_workers) as executor: - futures = {executor.submit(call_with_temp, q, t): i for i, (q, t) in enumerate(tasks)} - done = 0 - t0 = time.time() - for f in as_completed(futures): - done += 1 - r = f.result() - if r: - results.append(r) - if done % 50 == 0 or done == total: - elapsed = max(time.time() - t0, 1) - rate = len(results) / elapsed * 60 - print(f" 多温生成: {done}/{total} (+{len(results)}条, {rate:.0f}条/分)", end="\r") - - seen = set() - unique = [] - for item in results: - key = (item["question"], item["answer"]) - if key not in seen: - seen.add(key) - unique.append({"question": item["question"], "answer": item["answer"]}) - - elapsed = time.time() - t0 - rate = len(unique) / max(elapsed, 1) * 60 - print(f" 多温生成: 共{len(unique)}条 (去重后, {elapsed:.0f}s, {rate:.0f}条/分)") - return unique \ No newline at end of file diff --git a/build_vocab.py b/build_vocab.py deleted file mode 100644 index 92dc3ad..0000000 --- a/build_vocab.py +++ /dev/null @@ -1,134 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -用已有的 CharTokenizer + build_vocabulary 从训练数据生成词表 -然后直接启动训练 -""" -import sys, os, json, time -sys.path.insert(0, '/workspace/xiaoli-llm-chat') -sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1) -sys.stderr = os.fdopen(sys.stderr.fileno(), 'w', buffering=1) - -from llm_chat_improved import tokenizer_instance, build_vocabulary - -DATA_PATH = '/workspace/data/merged_clean_v3.json' -VOCAB_OUT = '/workspace/vocab.json' -VOCAB_OUT_XIAOLI = '/workspace/xiaoli-llm-chat/vocab.json' - -def load_pairs(path): - """加载训练数据为 (question, answer) 对""" - with open(path, 'r', encoding='utf-8') as f: - data = json.load(f) - - pairs = [] - for item in data: - if isinstance(item, dict): - q = item.get("question", "") or item.get("input", "") or item.get("prompt", "") - a = item.get("answer", "") or item.get("output", "") or item.get("response", "") - if q and a: - pairs.append((q, a)) - return pairs - -def main(): - print("=" * 60) - print("用 build_vocabulary 从训练数据生成词表") - print("=" * 60) - - # 1. 加载训练数据 - t0 = time.time() - print("\n[1] 加载训练数据...") - pairs = load_pairs(DATA_PATH) - print(f" 数据量: {len(pairs)} 对") - print(f" 耗时: {time.time()-t0:.1f}s") - - # 2. 用 build_vocabulary 生成词表 - t1 = time.time() - print("\n[2] 用 build_vocabulary 分析数据生成词表...") - vocab = build_vocabulary(pairs, min_freq=1) - - # build_vocabulary 返回的是 tokenizer_instance.char_to_id - # 已包含特殊token (=0, =1, etc.) - special_tokens = {'': 0, '': 1, '': 2, '': 3, - '': 4, '': 5, '': 6} - - real_vocab_size = max(vocab.values()) + 1 - - # 统计 - single_char_tokens = sum(1 for k in vocab if len(k) == 1) - multi_char_tokens = sum(1 for k in vocab if len(k) > 1 and not k.startswith('<')) - special_token_count = sum(1 for k in vocab if k.startswith('<')) - - print(f" 词表总项: {len(vocab)}") - print(f" vocab_size(模型用): {real_vocab_size}") - print(f" 单字符token: {single_char_tokens}") - print(f" 多字符token(英文词): {multi_char_tokens}") - print(f" 特殊token: {special_token_count}") - print(f" 耗时: {time.time()-t1:.1f}s") - - # 3. 编码测试 - print("\n[3] 编码测试:") - test_texts = [ - "你好呀！我是小狸～", - "小狸最喜欢吃草莓冰淇淋了！", - "你好我是小狸～", - ] - for text in test_texts: - ids = tokenizer_instance.encode(text) - decoded = tokenizer_instance.decode(ids) - print(f" 原文: '{text}'") - print(f" 编码: {ids[:20]}... (长度={len(ids)})") - print(f" 解码: '{decoded}'") - print() - - # 4. 估算模型参数量 - print("\n[4] 模型参数量估算 (d_model=320, n_layers=6, d_ff=1536):") - d_model = 320 - n_layers = 6 - d_ff = 1536 - vs = real_vocab_size - emb = vs * d_model - out = vs * d_model - per_layer = 4 * d_model**2 + 3 * d_model * d_ff + 2 * d_model - total = emb + out + n_layers * per_layer + 2 * d_model - print(f" vocab_size={vs}") - print(f" Embedding: {emb:,}") - print(f" fc_out: {out:,}") - print(f" Transformer({n_layers}层): {n_layers*per_layer:,}") - print(f" 总参数量: {total:,} ({total/1e6:.2f}M)") - - # 5. 保存词表 - print(f"\n[5] 保存词表...") - - # vocab.json 格式：不含特殊token，id从7开始（与 train.py 兼容） - normal_vocab = {k: v for k, v in vocab.items() if k not in special_tokens} - - # 备份旧词表 - for path in [VOCAB_OUT, VOCAB_OUT_XIAOLI]: - if os.path.exists(path): - backup = path.replace('vocab.json', 'vocab_old.json') - import shutil - shutil.copy2(path, backup) - print(f" 旧词表已备份: {backup}") - - # 保存到两个位置 - for path in [VOCAB_OUT, VOCAB_OUT_XIAOLI]: - with open(path, 'w', encoding='utf-8') as f: - json.dump(normal_vocab, f, ensure_ascii=False, indent=1) - print(f" {path}: {len(normal_vocab)}项 (不含特殊token, id从7开始)") - - # 保存完整词表（含特殊token，用于参考） - full_path = VOCAB_OUT_XIAOLI.replace('vocab.json', 'vocab_full.json') - with open(full_path, 'w', encoding='utf-8') as f: - json.dump(vocab, f, ensure_ascii=False, indent=1) - print(f" {full_path}: {len(vocab)}项 (含特殊token)") - - total_time = time.time() - t0 - print(f"\n总耗时: {total_time:.1f}s") - print("=" * 60) - print(f"词表生成完成！vocab_size = {real_vocab_size}") - print(f"模型参数量 ≈ {total/1e6:.2f}M") - print("=" * 60) - - -if __name__ == "__main__": - main() diff --git a/chat.html b/chat.html deleted file mode 100644 index db06c0f..0000000 --- a/chat.html +++ /dev/null @@ -1,284 +0,0 @@ - - - - - -小狸 v1.2 聊天 - - - -
-
-
🦊
-

小狸 v1.2

-

V2架构 · 21M参数 · 新训练启动 🔥

-
- -
- - -
- - - 0.15 - 🗑️ 清空 -
- -
- - -
-
按回车发送 · 建议温度 0.1~0.2
-
- - - - \ No newline at end of file diff --git a/chat_server.py b/chat_server.py deleted file mode 100644 index aa3a25e..0000000 --- a/chat_server.py +++ /dev/null @@ -1,143 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -小狸聊天 Web 服务 v2 - 支持多轮对话 -""" -import sys, os, json, re -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - -from flask import Flask, render_template, request, jsonify -import torch -from llm_chat_improved import ( - CharTokenizer, build_vocabulary, LLMModel, LLMChat, load_training_data -) - -app = Flask(__name__) - -# 全局模型 -llm_chat = None -vocab = None - -def clean_text(text): - """清洗文本""" - text = text.replace('\n', '') - text = re.sub(r'\*\*([^*]+)\*\*', r'\1', text) - text = re.sub(r'\([^)]*[ω｡ﾉ♡^][^)]*\)', '', text) - emoji_pattern = re.compile( - '[\U0001F600-\U0001F64F\U0001F300-\U0001F5FF\U0001F680-\U0001F6FF\U0001F1E0-\U0001F1FF\U00002702-\U000027B0\U0001f926-\U0001f937\u2640-\u2642\u2600-\u2B55\u200d\u23cf\u23e9\u231a\ufe0f\u3030\u2728\u2714\u2716\u27a1\u2b05\u2b06\u2b07\u2b1b\u2b1c\u2b50\u2b55\u2934\u2935\u25aa\u25ab\u25b6\u25c0\u25fb\u25fc\u25fd\u25fe]+', - flags=re.UNICODE - ) - text = emoji_pattern.sub('', text) - text = re.sub(r'[ \t]+', ' ', text) - return text.strip() - -def load_model(): - """加载训练好的模型""" - global llm_chat, vocab - - print("🐱 正在加载小狸模型...") - - device = 'cpu' - - # 直接从模型文件加载 vocab 和配置 (v4大模型) - checkpoint = torch.load('xiaoli_final_v4.pth', map_location=device) - vocab = checkpoint['vocab'] - config = checkpoint.get('model_config', {}) - - print(f" 词汇表大小: {len(vocab)}") - - # 创建模型 - 使用保存的配置（v4大模型带MoE+Thinking+Memory） - model = LLMModel( - vocab_size=config.get('vocab_size', len(vocab)), - d_model=config.get('d_model', 384), - n_heads=config.get('n_heads', 8), - n_layers=config.get('n_layers', 6), - d_ff=config.get('d_ff', 1536), - max_len=config.get('max_len', 128), - dropout=0.0, - use_memory=config.get('use_memory', True), - use_thinking=config.get('use_thinking', True), - use_moe=config.get('use_moe', True), - num_experts=config.get('num_experts', 4), - top_k=config.get('top_k', 2), - ) - - # 加载权重 - model.load_state_dict(checkpoint['model_state_dict']) - - llm_chat = LLMChat(model, vocab, device, 128) - print("✅ 小狸已就绪！") - -# 存储会话历史 -session_history = {} - -@app.route('/') -def index(): - return render_template('chat.html') - -@app.route('/chat', methods=['POST']) -def chat(): - data = request.json - message = data.get('message', '').strip() - session_id = data.get('session_id', 'default') - - if not message: - return jsonify({'response': '请输入内容哦~'}) - - if llm_chat is None: - return jsonify({'response': '模型加载中，请稍后再试'}) - - try: - # 获取会话历史 - history = session_history.get(session_id, []) - - # 构建带上下文的输入（使用与训练一致的格式） - # 训练格式: 问题 - # 多轮格式: 历史问题1历史回答1历史问题2历史回答2当前问题 - if history and len(history) >= 2: - # 只保留最近2轮对话（4条记录） - recent_history = history[-4:] - # 提取问题和答案 - context_parts = [] - for i in range(0, len(recent_history), 2): - if i + 1 < len(recent_history): - q = recent_history[i].replace("用户：", "") - a = recent_history[i+1].replace("小狸：", "") - context_parts.append(f"{q}{a}") - - if context_parts: - context = "".join(context_parts) - full_input = f"{context}{message}" - else: - full_input = message - else: - full_input = message - - # 限制输入长度（字符数） - if len(full_input) > 80: - full_input = full_input[-80:] - - response = llm_chat.generate_response(full_input, max_length=50, temperature=0.3) - - # 更新历史 - history.append(f"用户：{message}") - history.append(f"小狸：{response}") - # 只保留最近6轮（12条记录） - session_history[session_id] = history[-12:] - - return jsonify({'response': response}) - except Exception as e: - import traceback - return jsonify({'response': f'出错了: {str(e)}\n{traceback.format_exc()}'}) - -@app.route('/clear', methods=['POST']) -def clear_history(): - data = request.json - session_id = data.get('session_id', 'default') - session_history[session_id] = [] - return jsonify({'status': 'ok'}) - -if __name__ == '__main__': - load_model() - print("\n🌐 启动 Web 服务: http://0.0.0.0:5000") - app.run(host='0.0.0.0', port=5000, debug=False) diff --git a/distill_daemon.py b/distill_daemon.py deleted file mode 100644 index 912e0c5..0000000 --- a/distill_daemon.py +++ /dev/null @@ -1,197 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -小狸 AI - 持续蒸馏守护脚本 -============================ -每轮生成 1000 条新 QA 对，去重后追加到训练集。 -5 并发调用 Astron API，不超限。 - -使用方式: - python3 distill_daemon.py # 跑一轮 - python3 distill_daemon.py --loop # 持续循环，每轮间隔 60 秒 - python3 distill_daemon.py --count 5000 # 指定本轮目标条数 -""" -import sys, os, json, time, random, argparse -from concurrent.futures import ThreadPoolExecutor, as_completed -from pathlib import Path - -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -from api_correction import call_astron, batch_generate - -# ── 种子话题（15 大类，持续扩展） ────────────────────── -SEED_TOPICS = { - "日常问候": ["今天开心吗", "在干嘛呢", "吃了吗", "你好呀", "早上好"], - "情绪感受": ["我有点难过", "今天好开心", "我好无聊", "有点害怕", "好兴奋"], - "学校生活": ["上学好玩吗", "今天学了什么", "作业写完了吗", "喜欢上什么课", "考试考了多少"], - "朋友关系": ["你有好朋友吗", "怎么交朋友", "朋友吵架了", "一起玩什么", "好朋友搬家了"], - "家庭生活": ["喜欢爸爸妈妈吗", "有兄弟姐妹吗", "帮妈妈做家务", "周末去哪玩", "家里养宠物了吗"], - "兴趣爱好": ["喜欢唱歌吗", "会画画吗", "喜欢什么动画片", "爱看书吗", "喜欢什么游戏"], - "奇思妙想": ["为什么天是蓝的", "星星会掉下来吗", "鱼在水里怎么呼吸", "花为什么会开", "月亮上有兔子吗"], - "安慰鼓励": ["失败了怎么办", "被人欺负了", "好紧张怎么办", "觉得自己不够好", "想放弃了"], - "节日话题": ["最喜欢什么节日", "过年开心吗", "生日怎么过", "圣诞节收到什么礼物", "儿童节去哪玩"], - "自然科普": ["下雨了怎么办", "为什么会有彩虹", "树叶为什么变黄", "雪是怎么来的", "太阳下山去哪了"], - "童话故事": ["喜欢什么童话", "听过灰姑娘吗", "美人鱼的故事", "孙悟空厉害吗", "小红帽的故事"], - "美食话题": ["喜欢吃什么", "讨厌什么蔬菜", "冰淇淋好吃吗", "会做吃的吗", "妈妈做的饭好吃吗"], - "动物朋友": ["喜欢猫还是狗", "小兔子可爱吗", "小鸟会说话吗", "去过动物园吗", "恐龙是真的吗"], - "睡前话题": ["晚安怎么说", "做噩梦了", "睡不着怎么办", "想要抱抱", "讲个睡前故事"], - "成长烦恼": ["不想长大", "怕黑怎么办", "一个人在家害怕", "摔跤了很疼", "被批评了好难过"], -} - -# 所有话题平铺 -ALL_TOPICS = [] -for cat, items in SEED_TOPICS.items(): - ALL_TOPICS.extend(items) - -# 对话风格模板 -STYLE_TEMPLATES = [ - "请用8岁小女孩小狸的语气回答，简短自然，1-2句话。", - "小狸是个有点害羞但很可爱的小女孩，请模仿她的语气回答。", - "像小朋友一样简单回答就好，要用小狸的口吻哦。", - "用小狸的口吻回答，要可爱但不要太刻意，简短就好。", - "你是一个8岁小女孩，请用小孩子的方式回答这个问题。", -] - - -def generate_seed_questions(count: int = 500) -> list: - """从种子话题生成问题列表（随机组合+衍生）""" - questions = [] - - # 1. 直接使用种子话题 - for topic in random.sample(ALL_TOPICS, min(count // 3, len(ALL_TOPICS))): - questions.append(topic) - - # 2. 组合衍生（如 "今天开心吗 + 考试考了多少" → "今天考试考得开心吗"） - variant_templates = [ - lambda a, b: f"{a}，那{b}", - lambda a, b: f"如果{a}，是不是{b}", - lambda a, b: f"{a}？还有{b}", - lambda a, b: f"为什么{a}，然后{b}呀", - ] - needed = count // 3 - for i in range(needed): - a = random.choice(ALL_TOPICS) - b = random.choice(ALL_TOPICS) - tmpl = random.choice(variant_templates) - questions.append(tmpl(a, b)) - - # 3. "为什么"类衍生 - why_prefixes = ["为什么", "怎么才能", "可以教我", "告诉我", "什么是", "什么时候", "在哪里", "谁"] - for i in range(count - len(questions)): - base = random.choice(ALL_TOPICS) - prefix = random.choice(why_prefixes) - questions.append(f"{prefix}{base}?") - - random.shuffle(questions) - return questions[:count] - - -def load_existing_pairs(train_path: str) -> set: - """加载已有数据，供去重用""" - if not os.path.exists(train_path): - return set() - try: - with open(train_path, 'r', encoding='utf-8') as f: - data = json.load(f) - return set((item.get("question", ""), item.get("answer", "")) for item in data) - except: - return set() - - -def append_to_train(new_pairs: list, train_path: str, merged_path: str): - """去重后追加到训练集和合并文件""" - if not new_pairs: - return 0 - - existing = load_existing_pairs(train_path) - truly_new = [item for item in new_pairs if (item["question"], item["answer"]) not in existing] - - if not truly_new: - print(f" ℹ️ 无新内容（全部已存在）") - return 0 - - # 追加到训练集 - if os.path.exists(train_path): - with open(train_path, 'r', encoding='utf-8') as f: - cur_train = json.load(f) - else: - cur_train = [] - cur_train.extend(truly_new) - with open(train_path, 'w', encoding='utf-8') as f: - json.dump(cur_train, f, ensure_ascii=False, indent=2) - - # 同步到合并文件 - if os.path.exists(merged_path): - with open(merged_path, 'r', encoding='utf-8') as f: - cur_merged = json.load(f) - cur_merged.extend(truly_new) - with open(merged_path, 'w', encoding='utf-8') as f: - json.dump(cur_merged, f, ensure_ascii=False, indent=2) - - unique_qs = len(set(item["question"] for item in truly_new)) - print(f" ✅ 新增 {len(truly_new)} 条（{unique_qs} 道独特问题）到训练集") - return len(truly_new) - - -def run_distill_round(target_count: int = 1000, save_dir: str = None): - """执行一轮蒸馏：生成问题 → API 回答 → 去重追加""" - if save_dir is None: - save_dir = os.path.dirname(os.path.abspath(__file__)) - - train_path = os.path.join(save_dir, 'xiaoli_train_v2.json') - merged_path = os.path.join(save_dir, 'xiaoli_merged_data_v2.json') - - print(f"\n{'='*60}") - print(f" 蒸馏轮次开始 | 目标: {target_count} 条") - print(f"{'='*60}") - - # 1. 生成种子问题 - questions = generate_seed_questions(target_count) - print(f" 📝 生成 {len(questions)} 个问题") - - # 2. 用 Astron API 批量回答（5 并发） - t0 = time.time() - results = batch_generate(questions, temperature=0.85, max_workers=5) - elapsed = time.time() - t0 - - # 3. 去重追加 - added = append_to_train(results, train_path, merged_path) - - rate = added / max(elapsed, 1) * 60 - print(f" 📊 本轮: +{added} 条 | 耗时 {elapsed:.0f}s | 速率 {rate:.0f} 条/分") - print(f" 💾 训练集: {os.path.getsize(train_path)/1024/1024:.0f}MB (若存在)") - - return added - - -def main(): - parser = argparse.ArgumentParser(description="小狸 AI 持续蒸馏守护脚本") - parser.add_argument("--count", type=int, default=1000, help="每轮目标条数") - parser.add_argument("--loop", action="store_true", help="持续循环模式") - parser.add_argument("--interval", type=int, default=300, help="循环间隔秒数（默认300）") - parser.add_argument("--rounds", type=int, default=0, help="循环轮数（0=无限）") - args = parser.parse_args() - - if args.loop: - round_num = 1 - total_added = 0 - while True: - print(f"\n{'#'*60}") - print(f" # 第 {round_num} 轮蒸馏") - print(f"{'#'*60}") - added = run_distill_round(args.count) - total_added += added - print(f" 📈 累计新增: {total_added} 条") - - if args.rounds > 0 and round_num >= args.rounds: - print(f"\n✅ 完成 {args.rounds} 轮蒸馏，共新增 {total_added} 条") - break - - print(f" ⏳ 等待 {args.interval}s 后下一轮...") - time.sleep(args.interval) - round_num += 1 - else: - run_distill_round(args.count) - - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/gen_pretrain.py b/gen_pretrain.py deleted file mode 100644 index 64809d1..0000000 --- a/gen_pretrain.py +++ /dev/null @@ -1,484 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -小狸 AI - 预训练数据生成器 (v5) -================================== -两阶段生成： - -Phase 1 - 种子直出: - 5并发 API 直接生成完整多轮对话 - 每API调用 = 1条高质量对话 (15-25轮, 2000-4000字) - 速率: ~600条/h, 17h出1万条种子 - -Phase 2 - 语义扩展: - 基于种子主题，API重新"编剧"出新对话 - 每条种子 → 10个变体（换角度/深度/角色） - 速率: ~600条/h, ~7天扩展至10万+条 - -数据格式: JSONL, 每条: {"text":"...", "domain":"...", "concept":"..."} -""" -import sys, os, json, time, random, hashlib, argparse, re -from concurrent.futures import ThreadPoolExecutor, as_completed -from pathlib import Path - -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -from api_correction import call_astron - -# ═══════════════════════════════════════════════════════════ -# 1. 角色定义 -# ═══════════════════════════════════════════════════════════ - -TEACHER_SYSTEM = "你是一位AI导师，回答专业详细，包含具体数据和原理。每轮回答150-500字。" - -# ═══════════════════════════════════════════════════════════ -# 2. 深度学习领域图谱 -# ═══════════════════════════════════════════════════════════ - -DOMAINS = { - "量子物理": ["量子叠加态","量子纠缠","量子隧穿","海森堡测不准原理","波粒二象性","量子计算","量子场论","狄拉克方程"], - "相对论": ["狭义相对论","质能方程","时空弯曲","引力波","黑洞物理","虫洞理论","暗物质","暗能量"], - "人工智能": ["深度学习原理","Transformer","强化学习","大语言模型","扩散模型","多模态学习","注意力机制","模型量化"], - "计算机科学": ["计算理论","操作系统","分布式系统","数据库","编译器","密码学","网络协议","计算机视觉"], - "分子生物学": ["DNA复制","CRISPR","中心法则","蛋白质折叠","表观遗传学","RNA干扰","免疫系统","合成生物学"], - "数学": ["微积分","傅里叶分析","微分方程","抽象代数","数论","拓扑学","概率论","图论"], - "宇宙学": ["大爆炸","恒星演化","星系形成","宇宙微波背景","暗物质粒子","暴胀模型","引力波天文学","哈勃常数危机"], - "神经科学": ["神经元","突触可塑性","大脑皮层","记忆机制","睡眠与意识","视觉系统","脑机接口","情感神经科学"], - "化学": ["化学键","反应动力学","热力学","催化原理","光谱学","纳米化学","电化学","高分子化学"], - "工程科技": ["航天工程","半导体","通信工程","机器人学","电力系统","材料工程","机械设计","核工程"], - "经济学": ["一般均衡","博弈论","信息不对称","行为经济学","货币理论","计量经济学","拍卖理论","机制设计"], - "哲学": ["先验哲学","权力意志","存在主义","语言游戏","证伪主义","范式理论","正义论","功利主义"], - "天文学": ["赫罗图","超新星","中子星","系外行星","星系分类","脉冲星","恒星形成","宇宙射线"], - "地球科学": ["板块构造","地震学","火山学","古生物学","气候系统","海洋环流","大气动力学","行星地质"], - "生物技术": ["基因测序","PCR","冷冻电镜","核磁共振","光遗传学","类器官","纳米医学","基因敲除"], -} - -ALL_DOMAINS = list(DOMAINS.keys()) - - -# ═══════════════════════════════════════════════════════════ -# 3. Phase 1: API直接生成种子 -# ═══════════════════════════════════════════════════════════ - -def build_seed_prompt(domain: str, concept: str) -> str: - return f"""请围绕【{domain}】领域的【{concept}】概念，生成一段师生深度对话。 - -要求： -1. 学生：从基础概念问起，逐步深入，包含追问、质疑、要求举例 -2. 导师：每轮回答150-500字，专业准确，有数据支撑 -3. 对话演进：基础定义→核心机制→实验证据→应用场景→前沿进展 -4. 至少15轮对话 - -每轮格式（严格）： -学生：[问题] -导师：[回答] - -不要额外说明，直接开始对话。""" - - -def generate_seed(domain: str, concept: str) -> dict: - """一次API调用生成一条种子对话""" - prompt = build_seed_prompt(domain, concept) - text = call_astron(prompt, temperature=random.uniform(0.7, 0.9), - system_prompt=TEACHER_SYSTEM, max_tokens=8192) - if not text: - return None - - # 解析 - turns = [] - for line in text.strip().split('\n'): - line = line.strip() - if not line: - continue - if re.match(r'^学生[：:]', line): - content = re.sub(r'^学生[：:]\s*', '', line) - turns.append(('user', content)) - elif re.match(r'^导师[：:]', line): - content = re.sub(r'^导师[：:]\s*', '', line) - turns.append(('assistant', content)) - - if len(turns) < 10: - return None - - lines = [] - for role, content in turns: - lines.append(f"<|im_start|>{role}\n{content}<|im_end|>") - - return { - "text": '\n'.join(lines), - "domain": domain, - "concept": concept, - } - - -# ═══════════════════════════════════════════════════════════ -# 4. Phase 2: 语义扩展 -# ═══════════════════════════════════════════════════════════ - -VARIATION_PROMPTS = [ - "用更通俗的方式重新讲一遍，多用日常类比", - "从更高级的视角讲，假设学生已有基础，重点讲深层机制和数学原理", - "换个角度，从历史发展脉络讲，突出关键实验和突破性发现", - "从工程应用角度讲，重点讲实际怎么用、有什么技术瓶颈", - "从批判性视角讲，有哪些争议、局限性、未解决的问题", - "假设学生是高中水平，用最通俗的语言和大量类比来解释", - "从跨学科视角讲，这个概念与其他领域有什么联系", - "以问答竞赛的方式讲，学生提尖锐问题，导师用证据回击", -] - - -def generate_variation(seed: dict, seed_text: str, prompt_template: str) -> dict: - """基于种子主题生成一个语义变体""" - domain = seed["domain"] - concept = seed["concept"] - instruction = prompt_template - - prompt = f"""请围绕【{domain}】领域的【{concept}】概念，生成一段师生深度对话。 - -{instruction} - -要求： -1. 至少15轮对话 -2. 导师回答150-500字，专业详细 -3. 对话要有递进：定义→机制→证据→应用→前沿 - -每轮格式（严格）： -学生：[问题] -导师：[回答] - -直接开始对话，不要额外说明。""" - - text = call_astron(prompt, temperature=random.uniform(0.75, 0.95), - system_prompt=TEACHER_SYSTEM, max_tokens=8192) - if not text: - return None - - turns = [] - for line in text.strip().split('\n'): - line = line.strip() - if not line: - continue - if re.match(r'^学生[：:]', line): - content = re.sub(r'^学生[：:]\s*', '', line) - turns.append(('user', content)) - elif re.match(r'^导师[：:]', line): - content = re.sub(r'^导师[：:]\s*', '', line) - turns.append(('assistant', content)) - - if len(turns) < 10: - return None - - lines = [] - for role, content in turns: - lines.append(f"<|im_start|>{role}\n{content}<|im_end|>") - - return { - "text": '\n'.join(lines), - "domain": domain, - "concept": concept, - "variation": prompt_template[:30], - } - - -# ═══════════════════════════════════════════════════════════ -# 5. 数据管理 -# ═══════════════════════════════════════════════════════════ - -class DataStore: - def __init__(self, save_dir: str): - self.save_dir = Path(save_dir) - self.save_dir.mkdir(parents=True, exist_ok=True) - self.output_path = self.save_dir / 'dialogues.jsonl' - self.seed_path = self.save_dir / 'seeds.jsonl' - - self.existing_hashes = set() - self.count = 0 - self.load() - - def load(self): - for path in [self.output_path, self.seed_path]: - if path.exists(): - with open(path, 'r') as f: - for line in f: - line = line.strip() - if line: - try: - item = json.loads(line) - h = hashlib.md5(item['text'][:100].encode()).hexdigest() - self.existing_hashes.add(h) - self.count += 1 - except: - pass - print(f"📂 已有: {self.count:,} 条") - - def add_seeds(self, items: list) -> int: - if not items: - return 0 - new = [] - for item in items: - h = hashlib.md5(item['text'][:100].encode()).hexdigest() - if h not in self.existing_hashes: - self.existing_hashes.add(h) - new.append(item) - if not new: - return 0 - with open(self.seed_path, 'a') as f: - for item in new: - f.write(json.dumps(item, ensure_ascii=False) + '\n') - self.count += 1 - return len(new) - - def add_variations(self, items: list) -> int: - if not items: - return 0 - new = [] - for item in items: - h = hashlib.md5(item['text'][:100].encode()).hexdigest() - if h not in self.existing_hashes: - self.existing_hashes.add(h) - new.append(item) - if not new: - return 0 - with open(self.output_path, 'a') as f: - for item in new: - f.write(json.dumps(item, ensure_ascii=False) + '\n') - self.count += 1 - return len(new) - - def get_seeds(self): - """读取所有种子""" - seeds = [] - if self.seed_path.exists(): - with open(self.seed_path, 'r') as f: - for line in f: - line = line.strip() - if line: - seeds.append(json.loads(line)) - return seeds - - def total_count(self): - return self.count - - -# ═══════════════════════════════════════════════════════════ -# 6. 质量评分 -# ═══════════════════════════════════════════════════════════ - -def score(text: str) -> float: - turns = (text.count('<|im_start|>user') + text.count('<|im_start|>assistant')) // 2 - chars = len(text) - s = 0.0 - if turns >= 20: s += 0.20 - elif turns >= 12: s += 0.15 - elif turns >= 8: s += 0.10 - else: s += 0.05 - if chars >= 4000: s += 0.20 - elif chars >= 2000: s += 0.15 - elif chars >= 1200: s += 0.10 - else: s += 0.05 - avg = chars / max(turns, 1) - if avg >= 200: s += 0.15 - elif avg >= 100: s += 0.10 - uc = text.count('<|im_start|>user') - ac = text.count('<|im_start|>assistant') - ratio = min(uc, ac) / max(uc, ac, 1) - if ratio >= 0.8: s += 0.10 - unique = len(set(text)) - if unique >= 150: s += 0.10 - return max(0, min(1, s)) - - -# ═══════════════════════════════════════════════════════════ -# 7. 卷统计 -# ═══════════════════════════════════════════════════════════ - -def print_stats(store: DataStore, t_start: float, phase: str, target: int): - elapsed = max(time.time() - t_start, 1) - rate = store.count / elapsed * 3600 - eta_h = (target - store.count) / max(rate, 0.01) - - seed_count = 0 - if store.seed_path.exists(): - with open(store.seed_path) as f: - seed_count = sum(1 for _ in f if _.strip()) - - tokens_est = store.count * 6000 - - print(f" [{time.strftime('%H:%M:%S')}] " - f"Phase {phase} | " - f"累计: {store.count:,}/{target:,} " - f"({store.count/max(target,1)*100:.1f}%) | " - f"速率: {rate:.0f}条/h | " - f"种子: {seed_count} | " - f"tokens: {tokens_est/1e9:.2f}B | " - f"预计: {eta_h/24:.1f}d", flush=True) - - -# ═══════════════════════════════════════════════════════════ -# 8. 主函数 -# ═══════════════════════════════════════════════════════════ - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--target", type=int, default=100000, help="目标条数") - parser.add_argument("--phase", type=int, default=1, choices=[1, 2], - help="1=生成种子, 2=语义扩展") - parser.add_argument("--concurrency", type=int, default=5, help="API并发数") - args = parser.parse_args() - - save_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'pretrain_data') - store = DataStore(save_dir) - - target = args.target - if store.count >= target: - print(f"✅ 已达目标 {target:,} 条") - return - - print(f"\n{'='*50}") - print(f" Phase {args.phase}: {'种子生成' if args.phase == 1 else '语义扩展'}") - print(f" 目标: {target:,} 条") - print(f" 并发: {args.concurrency}") - print(f" 当前: {store.count:,}") - print(f"{'='*50}\n") - - t_start = time.time() - - if args.phase == 1: - # Phase 1: 纯API种子生成 - seeds = store.get_seeds() - # 从种子中的(domain, concept)对去重 - used = set() - for s in seeds: - used.add((s.get("domain",""), s.get("concept",""))) - print(f" 已有 {len(used)} 个不同(领域,概念)种子\n") - - while store.count < target: - batch_items = [] - tasks = [] - for _ in range(args.concurrency): - # 选一个未用过的(领域,概念) - for attempt in range(20): - d = random.choice(ALL_DOMAINS) - c = random.choice(DOMAINS[d]) - if (d, c) not in used or attempt > 15: - break - used.add((d, c)) - tasks.append((d, c)) - - with ThreadPoolExecutor(max_workers=args.concurrency) as ex: - futures = {ex.submit(generate_seed, d, c): (d, c) for d, c in tasks} - for f in as_completed(futures): - result = f.result() - if result and score(result["text"]) >= 0.5: - batch_items.append(result) - - if batch_items: - added = store.add_seeds(batch_items) - n = len(batch_items) - print(f" +{n}条种子 | 领域: {batch_items[0]['domain']}, " - f"概念: {batch_items[0]['concept']}...", flush=True) - - print_stats(store, t_start, "1", target) - - if store.count >= target: - break - - elif args.phase == 2: - # Phase 2: 语义扩展 - seeds = store.get_seeds() - if not seeds: - print("❌ 没有种子数据，先跑 Phase 1") - return - print(f" 种子库: {len(seeds)} 条\n") - - used_seeds = set() - expansion_tasks = [] - - while store.count < target: - if not expansion_tasks: - # 把所有种子打乱，生成扩展任务 - random.shuffle(seeds) - for s in seeds: - if s["text"] in used_seeds: - continue - used_seeds.add(s["text"]) - for vp in VARIATION_PROMPTS: - expansion_tasks.append((s, vp)) - print(f" 生成 {len(expansion_tasks)} 个扩展任务\n") - - # 取一批任务 - batch = [] - while len(batch) < args.concurrency and expansion_tasks: - batch.append(expansion_tasks.pop(0)) - - if not batch: - break - - with ThreadPoolExecutor(max_workers=args.concurrency) as ex: - futures = {ex.submit(generate_variation, s, s["text"], vp): (s, vp) - for s, vp in batch} - for f in as_completed(futures): - result = f.result() - if result and score(result["text"]) >= 0.5: - store.add_variations([result]) - - # 报告 - seed_count = 0 - if store.seed_path.exists(): - with open(store.seed_path) as f: - seed_count = sum(1 for _ in f if _.strip()) - elapsed = max(time.time() - t_start, 1) - rate = store.count / elapsed * 3600 - eta_h = (target - store.count) / max(rate, 0.01) if rate > 0 else 0 - tokens_est = store.count * 6000 - done_tasks = max(1, len(used_seeds) * len(VARIATION_PROMPTS) - len(expansion_tasks)) - print(f" [{time.strftime('%H:%M:%S')}] " - f"累计: {store.count:,}/{target:,} " - f"({store.count/max(target,1)*100:.1f}%) | " - f"速率: {rate:.0f}条/h | " - f"种子: {seed_count} | " - f"扩展任务进度: {done_tasks} | " - f"tokens: {tokens_est/1e9:.2f}B | " - f"预计: {eta_h/24:.1f}d", flush=True) - - if store.count >= target: - break - - # 最终报告 - total_time = time.time() - t_start - final_rate = store.count / max(total_time, 1) * 3600 - tokens = store.count * 6000 - path = store.output_path if args.phase == 2 else store.seed_path - size_mb = os.path.getsize(path) / 1024 / 1024 if path.exists() else 0 - - print(f"\n{'='*50}") - print(f" ✅ Phase {args.phase} 完成!") - print(f" 总条数: {store.count:,}") - print(f" 文件: {path}") - print(f" 大小: {size_mb:.1f}MB") - print(f" 估算tokens: {tokens/1e9:.2f}B") - print(f" 耗时: {total_time/3600:.1f}h") - print(f" 平均速率: {final_rate:.0f}条/h") - print(f"{'='*50}") - - # 样例 - print("\n📊 样例:") - with open(path) as f: - for i, line in enumerate(f): - if i >= 2: - break - if line.strip(): - item = json.loads(line.strip()) - text = item["text"] - turns = text.count("<|im_start|>") // 2 - print(f"\n 条目{i+1}: {turns}轮 | {len(text):,}字 | {item.get('domain','')}") - lines = text.split("\n")[:4] - for l in lines: - if l.strip(): - role = "👤" if "user" in l else "🧑‍🏫" - c = l.split("<|im_start|>")[1].split("<|im_end|>")[0].strip() - print(f" {role} {c[:100]}") - print(f" ...（{turns}轮）") - - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/llm_chat_improved.py b/llm_chat_improved.py deleted file mode 100644 index 62aec67..0000000 --- a/llm_chat_improved.py +++ /dev/null @@ -1,1330 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -LLM对话AI训练代码 - 改进版 -基于Transformer架构，参考MiniMind/nanoGPT等开源项目优化 -""" - -import torch -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -from torch.utils.data import Dataset, DataLoader -import math -import json -import os -from typing import List, Tuple, Dict, Any -import time -import threading -import queue -import re -from collections import defaultdict - - -# ==================== 1. 字符级分词器（对中文友好）==================== -class CharTokenizer: - """简单字符级分词器 - 中文按字符，英文按单词 - - 每个字符/单词作为一个token，对中文效果更好。 - """ - - def __init__(self): - self.special_tokens = { - '': 0, '': 1, '': 2, '': 3, - '': 4, '': 5, '': 6 - } - self.char_to_id = self.special_tokens.copy() - self.id_to_char = {v: k for k, v in self.char_to_id.items()} - self.next_id = len(self.special_tokens) - self._sorted_vocab_cache = None # encode缓存用 - - def build_vocab(self, texts: List[str], min_freq: int = 1) -> Dict[str, int]: - """从文本构建词汇表""" - char_counts = defaultdict(int) - - for text in texts: - # 中文按字符，英文按空格分割 - chars = self._split_chars(text) - for char in chars: - char_counts[char] += 1 - - # 添加高频字符 - for char, count in char_counts.items(): - if count >= min_freq and char not in self.char_to_id: - self.char_to_id[char] = self.next_id - self.id_to_char[self.next_id] = char - self.next_id += 1 - - # 清除encode缓存（vocab已变化） - self._sorted_vocab_cache = None - # 清除多字符token缓存 - for attr in ['_multi_trie', '_multi_keys', '_single_map']: - if hasattr(self, attr): - delattr(self, attr) - - return self.char_to_id - - def _split_chars(self, text: str) -> List[str]: - """分割文本为字符/词""" - # 先尝试按空格分割英文 - parts = text.split(' ') - result = [] - for part in parts: - if part.isascii(): - result.append(part) - else: - # 中文字符逐个添加 - for char in part: - if char.strip(): - result.append(char) - return result - - def encode(self, text: str) -> List[int]: - """将文本编码为ID序列 - 高效版""" - if not text: - return [self.special_tokens['']] - - ids = [] - i = 0 - n = len(text) - - # 构建多字符token的前缀映射（仅一次） - if not hasattr(self, '_multi_trie'): - self._multi_trie = {} - for token in self.char_to_id: - if len(token) > 1: - self._multi_trie[token] = self.char_to_id[token] - # 按长度降序排序，优先匹配长token - self._multi_keys = sorted(self._multi_trie.keys(), key=len, reverse=True) - self._single_map = {c: tid for c, tid in self.char_to_id.items() if len(c) == 1} - - while i < n: - matched = False - # 优先匹配多字符token（特殊标记） - for token in self._multi_keys: - if text.startswith(token, i): - ids.append(self._multi_trie[token]) - i += len(token) - matched = True - break - if not matched: - # 单字符直接用dict查 - c = text[i] - ids.append(self._single_map.get(c, self.special_tokens[''])) - i += 1 - - return ids - - def decode(self, ids: List[int]) -> str: - """将ID序列解码为文本""" - chars = [] - for tid in ids: - if tid in self.id_to_char: - char = self.id_to_char[tid] - # 保留特殊token以便后续处理，但这里先不过滤 - chars.append(char) - return ''.join(chars) - - -# 全局分词器实例 -tokenizer_instance = CharTokenizer() - - -def build_vocabulary(conversations: List[Tuple[str, str]], min_freq: int = 1) -> dict: - """构建词汇表""" - all_texts = [] - for question, answer in conversations: - all_texts.append(question) - all_texts.append(answer) - return tokenizer_instance.build_vocab(all_texts, min_freq) - - -# ==================== 2. 位置编码 ==================== -def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0): - """预计算旋转位置编码的频率""" - freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) - t = torch.arange(end, device=freqs.device) - freqs = torch.outer(t, freqs).float() - freqs_cos = torch.cos(freqs) - freqs_sin = torch.sin(freqs) - return freqs_cos, freqs_sin - - -def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor): - """调整频率张量形状以进行广播""" - ndim = x.ndim - assert 0 <= 1 < ndim - assert freqs_cis.shape == (x.shape[1], x.shape[-1]) - shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] - return freqs_cis.view(shape) - - -def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor): - """应用旋转位置编码 - 参考MiniMind实现""" - # 将xq和xk的最后一个维度分成两半 - xq_r, xq_i = xq.float().reshape(xq.shape[:-1] + (-1, 2)).unbind(-1) - xk_r, xk_i = xk.float().reshape(xk.shape[:-1] + (-1, 2)).unbind(-1) - - # 调整freqs形状 - freqs_cos = reshape_for_broadcast(freqs_cos, xq_r) - freqs_sin = reshape_for_broadcast(freqs_sin, xq_r) - - # 应用旋转 - xq_out_r = xq_r * freqs_cos - xq_i * freqs_sin - xq_out_i = xq_r * freqs_sin + xq_i * freqs_cos - xk_out_r = xk_r * freqs_cos - xk_i * freqs_sin - xk_out_i = xk_r * freqs_sin + xk_i * freqs_cos - - # 重新组合 - xq_out = torch.stack([xq_out_r, xq_out_i], dim=-1).flatten(3) - xk_out = torch.stack([xk_out_r, xk_out_i], dim=-1).flatten(3) - - return xq_out.type_as(xq), xk_out.type_as(xk) - - -# ==================== 3. RMSNorm ==================== -class RMSNorm(nn.Module): - """RMSNorm归一化层""" - - def __init__(self, dim: int, eps: float = 1e-6): - super().__init__() - self.eps = eps - self.weight = nn.Parameter(torch.ones(dim)) - - def forward(self, x): - rms = torch.sqrt(torch.mean(x.pow(2), dim=-1, keepdim=True) + self.eps) - return self.weight * x / rms - - -# ==================== 4. 注意力机制 ==================== -class Attention(nn.Module): - """多头注意力 - 参考MiniMind优化实现""" - - def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1): - super().__init__() - assert d_model % n_heads == 0 - self.d_model = d_model - self.n_heads = n_heads - self.d_head = d_model // n_heads - - # QKV投影合并为一个矩阵（更高效） - self.wq = nn.Linear(d_model, d_model, bias=False) - self.wk = nn.Linear(d_model, d_model, bias=False) - self.wv = nn.Linear(d_model, d_model, bias=False) - self.wo = nn.Linear(d_model, d_model, bias=False) - - self.dropout = nn.Dropout(dropout) - - def forward(self, x, freqs_cos, freqs_sin, mask=None): - batch_size, seq_len, _ = x.shape - - # QKV投影 - xq = self.wq(x) - xk = self.wk(x) - xv = self.wv(x) - - # 重塑为多头 [batch, seq, n_heads, d_head] - xq = xq.view(batch_size, seq_len, self.n_heads, self.d_head) - xk = xk.view(batch_size, seq_len, self.n_heads, self.d_head) - xv = xv.view(batch_size, seq_len, self.n_heads, self.d_head) - - # 应用RoPE - xq, xk = apply_rotary_emb(xq, xk, freqs_cos, freqs_sin) - - # 转置为 [batch, n_heads, seq, d_head] - xq = xq.transpose(1, 2) - xk = xk.transpose(1, 2) - xv = xv.transpose(1, 2) - - # 融合注意力（SDPA / FlashAttention-2）- 纯因果掩码，等价于原下三角 mask - dropout_p = self.dropout.p if self.training else 0.0 - output = F.scaled_dot_product_attention( - xq, xk, xv, attn_mask=None, dropout_p=dropout_p, is_causal=True) - - output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) - - return self.wo(output) - - -# ==================== 5. 前馈网络 ==================== -class FeedForward(nn.Module): - """SwiGLU前馈网络 - 现代LLM标准实现""" - - def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1): - super().__init__() - self.w1 = nn.Linear(d_model, d_ff, bias=False) - self.w2 = nn.Linear(d_ff, d_model, bias=False) - self.w3 = nn.Linear(d_model, d_ff, bias=False) - self.dropout = nn.Dropout(dropout) - - def forward(self, x): - return self.w2(self.dropout(nn.functional.silu(self.w1(x)) * self.w3(x))) - - -# ==================== 6. MoE专家混合 ==================== -class FeedForwardExpert(nn.Module): - """MoE专家 - 单个前馈网络""" - - def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1): - super().__init__() - self.w1 = nn.Linear(d_model, d_ff, bias=False) - self.w2 = nn.Linear(d_ff, d_model, bias=False) - self.w3 = nn.Linear(d_model, d_ff, bias=False) - self.dropout = nn.Dropout(dropout) - - def forward(self, x): - return self.w2(self.dropout(nn.functional.silu(self.w1(x)) * self.w3(x))) - - -class MixtureOfExperts(nn.Module): - """专家混合(MoE)层 - 高效并行实现""" - - def __init__(self, d_model: int, d_ff: int, num_experts: int = 4, top_k: int = 2, dropout: float = 0.1): - super().__init__() - self.num_experts = num_experts - self.top_k = top_k - self.d_model = d_model - - # 专家网络 - self.experts = nn.ModuleList([ - FeedForwardExpert(d_model, d_ff, dropout) for _ in range(num_experts) - ]) - - # 门控网络 - self.gate = nn.Linear(d_model, num_experts, bias=False) - - def forward(self, x): - batch_size, seq_len, d_model = x.shape - x_flat = x.view(-1, d_model) # [batch*seq, d_model] - - # 计算门控权重 - gate_logits = self.gate(x_flat) # [batch*seq, num_experts] - - # 选择top-k专家 - top_k_logits, top_k_indices = torch.topk(gate_logits, self.top_k, dim=-1) - top_k_weights = torch.softmax(top_k_logits, dim=-1) # [batch*seq, top_k] - - # 高效并行计算：为每个专家收集输入 - output = torch.zeros_like(x_flat) - - for expert_idx in range(self.num_experts): - # 找到选择该专家的所有位置 - mask = (top_k_indices == expert_idx).any(dim=-1) # [batch*seq] - if not mask.any(): - continue - - # 获取选择该专家的输入 - expert_input = x_flat[mask] # [num_selected, d_model] - - # 计算该专家的权重 - expert_positions = mask.nonzero(as_tuple=True)[0] - expert_weights = [] - for pos in expert_positions: - # 找到该位置在top_k中的权重 - pos_in_topk = (top_k_indices[pos] == expert_idx).nonzero(as_tuple=True)[0] - if len(pos_in_topk) > 0: - expert_weights.append(top_k_weights[pos, pos_in_topk[0]]) - else: - expert_weights.append(torch.tensor(0.0, device=x.device)) - expert_weights = torch.stack(expert_weights).unsqueeze(-1) # [num_selected, 1] - - # 前向传播 - expert_output = self.experts[expert_idx](expert_input) * expert_weights - - # 累加到输出 - output[mask] += expert_output - - return output.view(batch_size, seq_len, d_model) - - -# ==================== 7. Transformer块 ==================== -class TransformerBlock(nn.Module): - """Transformer块 - Pre-Norm + 可选MoE""" - - def __init__(self, d_model: int, n_heads: int, d_ff: int, dropout: float = 0.1, - use_moe: bool = False, num_experts: int = 4, top_k: int = 2): - super().__init__() - - self.attention = Attention(d_model, n_heads, dropout) - - # Pre-Norm - self.attn_norm = RMSNorm(d_model) - self.ff_norm = RMSNorm(d_model) - - # 前馈网络（可选MoE） - if use_moe: - self.feed_forward = MixtureOfExperts(d_model, d_ff, num_experts, top_k, dropout) - else: - self.feed_forward = FeedForward(d_model, d_ff, dropout) - - self.dropout = nn.Dropout(dropout) - - def forward(self, x, freqs_cos, freqs_sin, mask=None): - # Pre-Norm + 注意力 + 残差 - h = x + self.dropout(self.attention(self.attn_norm(x), freqs_cos, freqs_sin, mask)) - # Pre-Norm + 前馈 + 残差 - out = h + self.dropout(self.feed_forward(self.ff_norm(h))) - return out - - -# ==================== 8. 简化版思维模块 ==================== -class ThinkingModule(nn.Module): - """轻量级推理增强模块 - 通过中间层表示增强""" - - def __init__(self, d_model: int, n_heads: int, n_thinking_steps: int = 3): - super().__init__() - self.n_thinking_steps = n_thinking_steps - - # 多层推理处理 - self.thinking_layers = nn.ModuleList([ - nn.TransformerEncoderLayer( - d_model=d_model, - nhead=n_heads, - dim_feedforward=d_model * 4, - dropout=0.1, - batch_first=True - ) for _ in range(n_thinking_steps) - ]) - - self.norms = nn.ModuleList([RMSNorm(d_model) for _ in range(n_thinking_steps)]) - self.output_proj = nn.Linear(d_model, d_model) - - def forward(self, x): - current = x - for i in range(self.n_thinking_steps): - # 因果掩码 - seq_len = x.size(1) - mask = torch.triu(torch.ones(seq_len, seq_len, device=x.device), diagonal=1).bool() - attended = self.thinking_layers[i](current, src_mask=mask) - current = self.norms[i](current + attended) - - return self.output_proj(current) - - -# ==================== 9. 记忆模块 ==================== -class MemoryModule(nn.Module): - """记忆模块 - 简化版可学习记忆""" - - def __init__(self, d_model: int, memory_size: int = 10): - super().__init__() - self.d_model = d_model - self.memory_size = memory_size - - # 可学习的记忆键值对 - self.memory_keys = nn.Parameter(torch.randn(memory_size, d_model) * 0.02) - self.memory_values = nn.Parameter(torch.randn(memory_size, d_model) * 0.02) - - def forward(self, query): - """检索记忆""" - batch_size = query.size(0) - keys = self.memory_keys.unsqueeze(0).expand(batch_size, -1, -1) - values = self.memory_values.unsqueeze(0).expand(batch_size, -1, -1) - - # 计算注意力分数 - scores = torch.matmul(query, keys.transpose(-2, -1)) / math.sqrt(self.d_model) - attn_weights = torch.softmax(scores, dim=-1) - - memory_output = torch.matmul(attn_weights, values) - return memory_output - - -# ==================== 10. 主模型 ==================== -class LLMModel(nn.Module): - """改进版LLM模型""" - - def __init__(self, vocab_size: int, d_model: int = 512, n_heads: int = 8, - n_layers: int = 6, d_ff: int = 2048, max_len: int = 512, dropout: float = 0.1, - use_memory: bool = True, memory_size: int = 10, - use_thinking: bool = True, n_thinking_steps: int = 3, - use_moe: bool = True, num_experts: int = 4, top_k: int = 2): - super().__init__() - - self.vocab_size = vocab_size - self.d_model = d_model - self.n_heads = n_heads - self.n_layers = n_layers - self.d_ff = d_ff - self.max_len = max_len - self.use_memory = use_memory - self.use_thinking = use_thinking - self.use_moe = use_moe - - # 词嵌入 - self.embedding = nn.Embedding(vocab_size, d_model) - - # 预计算RoPE频率 - freqs_cos, freqs_sin = precompute_freqs_cis(d_model // n_heads, max_len * 2) - self.register_buffer('freqs_cos', freqs_cos) - self.register_buffer('freqs_sin', freqs_sin) - - # 思考模块 - if use_thinking: - self.thinking = ThinkingModule(d_model, n_heads, n_thinking_steps) - - # 记忆模块 - if use_memory: - self.memory = MemoryModule(d_model, memory_size) - - # Transformer层 - self.transformer_blocks = nn.ModuleList([ - TransformerBlock(d_model, n_heads, d_ff, dropout, - use_moe=use_moe, num_experts=num_experts, top_k=top_k) - for _ in range(n_layers) - ]) - - # 输出层 - self.fc_out = nn.Linear(d_model, vocab_size, bias=False) - - self.dropout = nn.Dropout(dropout) - - # 初始化 - self._init_weights() - - def _init_weights(self): - """参数初始化""" - for name, param in self.named_parameters(): - if param.dim() > 1: - if 'embedding' in name or 'fc_out' in name: - torch.nn.init.normal_(param, mean=0.0, std=0.02) - else: - torch.nn.init.normal_(param, mean=0.0, std=0.02) - - def forward(self, x, mask=None, return_thinking=False): - batch_size, seq_len = x.shape - - # 词嵌入 - x = self.embedding(x) - x = self.dropout(x) - - # 获取位置编码 - freqs_cos = self.freqs_cos[:seq_len] - freqs_sin = self.freqs_sin[:seq_len] - - thinking_output = None - - # 思考模块 - if self.use_thinking: - thinking_output = self.thinking(x) - x = x + thinking_output # 残差融合 - - # 记忆模块 - if self.use_memory: - memory_output = self.memory(x) - x = x + memory_output # 残差融合 - - # Transformer层 - for block in self.transformer_blocks: - x = block(x, freqs_cos, freqs_sin, mask) - - # 输出 - output = self.fc_out(x) - - if return_thinking and thinking_output is not None: - return output, thinking_output - return output - - def generate_mask(self, seq_len): - """生成因果掩码""" - mask = torch.tril(torch.ones(seq_len, seq_len)) - return mask.unsqueeze(0).unsqueeze(0) # [1, 1, seq, seq] - - -# ==================== 11. 数据集 ==================== -class ConversationDataset(Dataset): - """对话数据集""" - - def __init__(self, conversations: List[Tuple[str, str]], vocab: dict, max_length: int = 256): - self.conversations = conversations - self.vocab = vocab - self.max_length = max_length - self.pad_token_id = vocab.get('', 0) - self.unk_token_id = vocab.get('', 1) - - def __len__(self): - return len(self.conversations) - - def __getitem__(self, idx): - question, answer = self.conversations[idx] - - # 创建训练样本: 问题 答案 - text = f"{question}{answer}" - ids = tokenizer_instance.encode(text) - - # 填充或截断 - ids = self.pad_or_truncate(ids) - - # 创建输入和目标（错位一个位置） - input_ids = ids[:-1] - target_ids = ids[1:] - - return { - 'input_ids': torch.tensor(input_ids, dtype=torch.long), - 'target_ids': torch.tensor(target_ids, dtype=torch.long) - } - - def pad_or_truncate(self, ids: List[int]) -> List[int]: - if len(ids) > self.max_length: - return ids[:self.max_length] - else: - return ids + [self.pad_token_id] * (self.max_length - len(ids)) - - -# ==================== 12. 训练器 ==================== -class LLMTrainer: - """LLM训练器""" - - def __init__(self, model: nn.Module, device: str = 'cpu', lr: float = 0.0001, use_amp: bool = None): - self.model = model.to(device) - self.device = device - # 火力全开：TF32 加速 fp32 矩阵乘法（L40/A100/H100 原生支持） - torch.backends.cuda.matmul.allow_tf32 = True - torch.backends.cudnn.allow_tf32 = True - self.criterion = nn.CrossEntropyLoss(ignore_index=0) - self.optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01) - self.scheduler = optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=100) - - # 混合精度训练 - if use_amp is None: - self.use_amp = torch.cuda.is_available() - else: - self.use_amp = use_amp and torch.cuda.is_available() - - if self.use_amp: - self.scaler = torch.cuda.amp.GradScaler() - print("✓ 已启用混合精度训练（AMP）") - else: - self.scaler = None - - if torch.cuda.is_available(): - torch.backends.cudnn.benchmark = True - torch.backends.cudnn.deterministic = False - - self.save_queue = queue.Queue() - self.save_lock = threading.Lock() - - def train_epoch(self, dataloader: DataLoader, epoch: int, vocab: dict = None, config: dict = None): - self.model.train() - total_loss = 0 - accumulation_steps = config.get('gradient_accumulation_steps', 1) if config else 1 - log_interval = max(1, len(dataloader) // 100) # 每1%打印一次 - - for batch_idx, batch in enumerate(dataloader): - input_ids = batch['input_ids'].to(self.device, non_blocking=True) - target_ids = batch['target_ids'].to(self.device, non_blocking=True) - - # 生成因果掩码 - mask = self.model.generate_mask(input_ids.size(1)).to(self.device) - - # CUDA graphs step 边界标记（torch.compile reduce-overhead 必需，避免输出 buffer 被覆盖） - torch.compiler.cudagraph_mark_step_begin() - # 前向传播 - if self.use_amp: - with torch.cuda.amp.autocast(): - outputs = self.model(input_ids, mask) - outputs = outputs.reshape(-1, outputs.size(-1)) - target_ids = target_ids.reshape(-1) - loss = self.criterion(outputs, target_ids) - else: - outputs = self.model(input_ids, mask) - outputs = outputs.reshape(-1, outputs.size(-1)) - target_ids = target_ids.reshape(-1) - loss = self.criterion(outputs, target_ids) - - loss = loss / accumulation_steps - - # 反向传播 - if self.use_amp: - self.scaler.scale(loss).backward() - if (batch_idx + 1) % accumulation_steps == 0: - self.scaler.unscale_(self.optimizer) - torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) - self.scaler.step(self.optimizer) - self.scaler.update() - self.optimizer.zero_grad() - else: - loss.backward() - if (batch_idx + 1) % accumulation_steps == 0: - torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) - self.optimizer.step() - self.optimizer.zero_grad() - - total_loss += loss.item() * accumulation_steps - - # 定期打印日志 - if (batch_idx + 1) % log_interval == 0 or batch_idx == len(dataloader) - 1: - avg_loss = total_loss / (batch_idx + 1) - mem_info = "" - if torch.cuda.is_available(): - allocated = torch.cuda.memory_allocated(self.device) / 1e9 - reserved = torch.cuda.memory_reserved(self.device) / 1e9 - mem_info = f", Mem: {allocated:.2f}GB/{reserved:.2f}GB" - print(f'Epoch [{epoch}], Batch [{batch_idx + 1}/{len(dataloader)}], Loss: {avg_loss:.4f}{mem_info}') - - self.scheduler.step() - return total_loss / len(dataloader) - - def evaluate(self, dataloader: DataLoader, config: dict = None): - """在验证集上计算loss（无梯度，不更新权重）""" - self.model.eval() - total_loss = 0 - accumulation_steps = config.get('gradient_accumulation_steps', 1) if config else 1 - - with torch.no_grad(): - for batch_idx, batch in enumerate(dataloader): - input_ids = batch['input_ids'].to(self.device, non_blocking=True) - target_ids = batch['target_ids'].to(self.device, non_blocking=True) - - mask = self.model.generate_mask(input_ids.size(1)).to(self.device) - torch.compiler.cudagraph_mark_step_begin() - outputs = self.model(input_ids, mask) - outputs = outputs.reshape(-1, outputs.size(-1)) - tgt = target_ids.reshape(-1) - - loss = self.criterion(outputs, tgt) - total_loss += loss.item() - - if len(dataloader) == 0: - return float('inf') - return total_loss / len(dataloader) - - def save_model(self, filepath: str, vocab: dict, epoch: int = None): - # 兼容 torch.compile：取原始（未编译）模块，避免 state_dict 带 _orig_mod. 前缀 - core = getattr(self.model, '_orig_mod', self.model) - save_dict = { - 'model_state_dict': core.state_dict(), - 'optimizer_state_dict': self.optimizer.state_dict(), - 'vocab': vocab, - 'model_config': { - 'vocab_size': core.vocab_size, - 'd_model': core.d_model, - 'n_heads': core.n_heads, - 'n_layers': core.n_layers, - 'd_ff': core.d_ff, - 'max_len': core.max_len, - 'use_memory': core.use_memory, - 'use_thinking': core.use_thinking, - 'use_moe': core.use_moe - } - } - if epoch is not None: - save_dict['epoch'] = epoch - - torch.save(save_dict, filepath) - print(f"模型已保存到: {filepath}") - - def load_model(self, filepath: str): - checkpoint = torch.load(filepath, map_location=self.device) - core = getattr(self.model, '_orig_mod', self.model) - core.load_state_dict(checkpoint['model_state_dict']) - self.optimizer.load_state_dict(checkpoint['optimizer_state_dict']) - vocab = checkpoint['vocab'] - print(f"模型已从 {filepath} 加载") - return vocab - - def request_save(self, filename: str, vocab: dict): - """请求在训练过程中保存模型""" - self.save_queue.put({'filename': filename, 'vocab': vocab}) - - -# ==================== 13. 对话生成 ==================== -class LLMChat: - """LLM对话""" - - def __init__(self, model: nn.Module, vocab: dict, device: str = 'cpu', max_length: int = 256, - show_thinking: bool = True): - self.model = model.to(device) - self.model.eval() - self.vocab = vocab - self.device = device - self.max_length = max_length - self.id_to_token = {v: k for k, v in vocab.items()} - self.show_thinking = show_thinking - self.conversation_history = [] - self.max_history = 10 - - # 同步全局tokenizer的vocab（关键！） - tokenizer_instance.char_to_id = vocab - tokenizer_instance.id_to_char = {v: k for k, v in vocab.items()} - # 清除encode缓存，避免用旧的cache导致错位 - tokenizer_instance._sorted_vocab_cache = None - for attr in ['_multi_trie', '_multi_keys', '_single_map']: - if hasattr(tokenizer_instance, attr): - delattr(tokenizer_instance, attr) - - def generate_response(self, question: str, max_length: int = 100, temperature: float = 0.8, - include_thinking: bool = False, use_beam_search: bool = False, - history: list = None): - """生成回答 - history: 可选，历史对话列表 [(user_q, assistant_a), ...] - """ - self.model.eval() - - # 构建输入（支持多轮上下文） - if history: - parts = [] - for hq, ha in history: - parts.append(hq) - parts.append(ha) - parts.append(question) - input_text = f"{''.join(parts)}" - else: - input_text = f"{question}" - - input_ids = tokenizer_instance.encode(input_text) - - if use_beam_search: - return self._generate_with_beam_search(input_ids, max_length, temperature, include_thinking) - else: - return self._generate_with_sampling(input_ids, max_length, temperature, include_thinking) - - def _generate_with_sampling(self, input_ids: List[int], max_length: int, temperature: float, include_thinking: bool): - """使用采样生成文本，带重复惩罚""" - input_tensor = torch.tensor([input_ids], dtype=torch.long).to(self.device) - generated = input_ids.copy() - - with torch.no_grad(): - for step in range(max_length): - seq = input_tensor[:, -self.model.max_len:] - - # 前向传播 - outputs = self.model(seq) - - # 获取最后一个时间步的预测 - logits = outputs[0, -1, :] - - # 温度=0 → greedy解码（取概率最高的token） - if temperature == 0: - next_token = torch.argmax(logits).item() - generated.append(next_token) - new_token = torch.tensor([[next_token]], device=self.device) - input_tensor = torch.cat([input_tensor, new_token], dim=1) - continue - - next_token_logits = logits / temperature - - # === 重复惩罚 === - # 1. 对最近生成的token进行惩罚（更强） - recent_tokens = generated[-30:] # 最近30个token - for token_id in set(recent_tokens): - count = recent_tokens.count(token_id) - if count > 1: - next_token_logits[token_id] -= 3.0 * count - - # 2. 阻止2-gram重复（更强） - if len(generated) >= 2: - last_bigram = (generated[-2], generated[-1]) - for i in range(len(generated) - 2): - if (generated[i], generated[i+1]) == last_bigram: - if i + 2 < len(generated): - next_token_logits[generated[i+2]] -= 10.0 - - # 3. 阻止3-gram重复 - if len(generated) >= 3: - last_trigram = (generated[-3], generated[-2], generated[-1]) - for i in range(len(generated) - 3): - if (generated[i], generated[i+1], generated[i+2]) == last_trigram: - if i + 3 < len(generated): - next_token_logits[generated[i+3]] -= 15.0 - - # 4. 阻止完全重复句子（检测连续重复模式） - decoded_so_far = tokenizer_instance.decode(generated[len(input_ids):]) - if self._has_repeated_pattern(decoded_so_far): - break - - # Top-k采样 - next_token = self._top_k_sampling(next_token_logits, top_k=30) - - # 检查结束标记 - if next_token == self.vocab.get('', -1): - break - - generated.append(next_token) - - # 更新输入 - new_token = torch.tensor([[next_token]], device=self.device) - input_tensor = torch.cat([input_tensor, new_token], dim=1) - - if len(generated) > max_length: - break - - # 解码 - generated_text = tokenizer_instance.decode(generated[len(input_ids):]) - generated_text = self._post_process_response(generated_text) - - return generated_text - - def _has_repeated_pattern(self, text: str) -> bool: - """检测文本中是否有重复3次以上的模式""" - if len(text) < 6: - return False - # 检测连续重复2次以上的子串（长度>=2） - for length in range(2, min(12, len(text)//2 + 1)): - for i in range(len(text) - 2*length + 1): - pattern = text[i:i+length] - if text[i:i+2*length] == pattern * 2: - return True - return False - - def _generate_with_beam_search(self, input_ids: List[int], max_length: int, temperature: float, - include_thinking: bool, beam_width: int = 3): - """使用束搜索生成文本""" - beams = [(input_ids.copy(), 0.0)] - - for _ in range(max_length): - candidates = [] - - for seq, score in beams: - input_tensor = torch.tensor([seq], dtype=torch.long).to(self.device) - - with torch.no_grad(): - outputs = self.model(input_tensor[:, -self.model.max_len:]) - next_token_logits = outputs[0, -1, :] / temperature - next_token_probs = torch.softmax(next_token_logits, dim=-1) - - top_k_probs, top_k_indices = torch.topk(next_token_probs, beam_width) - - for i in range(beam_width): - token_id = top_k_indices[i].item() - token_prob = top_k_probs[i].item() - new_seq = seq + [token_id] - new_score = score + math.log(token_prob) - candidates.append((new_seq, new_score)) - - candidates.sort(key=lambda x: x[1], reverse=True) - beams = candidates[:beam_width] - - # 检查是否结束 - if all(beam[0][-1] == self.vocab.get('', -1) for beam in beams): - break - - best_seq = beams[0][0] - generated_text = tokenizer_instance.decode(best_seq[len(input_ids):]) - generated_text = self._post_process_response(generated_text) - - return generated_text - - def _top_k_sampling(self, logits, top_k=50): - """Top-k采样""" - if top_k > 0: - top_k = min(top_k, logits.size(-1)) - indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None] - logits = logits.masked_fill(indices_to_remove, float('-inf')) - - probs = torch.softmax(logits, dim=-1) - next_token = torch.multinomial(probs, num_samples=1) - return next_token.item() - - def _post_process_response(self, response: str) -> str: - """后处理响应""" - response = response.replace('', '').replace('', '') - response = response.replace('', '').replace('', '') - response = response.replace('', '') - response = response.strip() - - # 去除重复片段（如 "我是小狸，我是小狸"） - response = self._remove_repeated_phrases(response) - - # 确保句子完整性 - if response and response[-1] not in ['。', '！', '？', '.', '!', '?', '\n']: - last_punct = max(response.rfind(p) for p in ['。', '！', '？', '.', '!', '?', '\n']) - if last_punct != -1: - response = response[:last_punct+1] - - return response if response else '...' - - def _remove_repeated_phrases(self, text: str) -> str: - """去除文本中的重复短语""" - if len(text) < 4: - return text - - # 检测并去除重复2次以上的子串 - for length in range(4, min(20, len(text)//2 + 1)): - i = 0 - while i <= len(text) - 2*length: - pattern = text[i:i+length] - if text[i:i+2*length] == pattern * 2: - # 找到重复，保留一个 - text = text[:i] + pattern + text[i+2*length:] - # 继续检查同一位置 - else: - i += 1 - - # 清理多余的标点 - text = re.sub(r'([。！？.!?])\1+', r'\1', text) - - return text.strip() - - def clear_history(self): - """清除对话历史""" - self.conversation_history = [] - - def get_history(self): - """获取对话历史""" - return self.conversation_history - - def toggle_thinking(self): - """切换是否显示思考过程""" - self.show_thinking = not self.show_thinking - return self.show_thinking - - -# ==================== 14. 数据加载 ==================== -def load_training_data(data_path: str) -> List[Tuple[str, str]]: - """加载训练数据""" - possible_paths = [ - data_path, - os.path.join(os.path.dirname(__file__), data_path), - os.path.join(os.getcwd(), data_path) - ] - - for path in possible_paths: - try: - if os.path.exists(path): - with open(path, 'r', encoding='utf-8') as f: - data = json.load(f) - - if isinstance(data, list): - if data and isinstance(data[0], dict): - if 'question' in data[0] and 'answer' in data[0]: - return [(item['question'], item['answer']) for item in data] - elif 'input' in data[0] and 'output' in data[0]: - return [(item['input'], item['output']) for item in data] - elif 'prompt' in data[0] and 'response' in data[0]: - return [(item['prompt'], item['response']) for item in data] - elif 'instruction' in data[0] and 'output' in data[0]: - return [(item['instruction'], item['output']) for item in data] - elif 'conversations' in data[0]: - conversations = [] - for item in data: - convs = item['conversations'] - for i in range(0, len(convs)-1, 2): - if i+1 < len(convs): - conversations.append((convs[i]['value'], convs[i+1]['value'])) - return conversations - elif 'human' in data[0] and 'assistant' in data[0]: - return [(item['human'], item['assistant']) for item in data] - else: - print(f"未知的数据格式，字段包括: {list(data[0].keys())}") - return [] - elif data and isinstance(data[0], list): - return [(str(item[0]), str(item[1])) for item in data if len(item) >= 2] - else: - return [] - elif isinstance(data, dict): - if 'data' in data and isinstance(data['data'], list): - if data['data'] and 'question' in data['data'][0] and 'answer' in data['data'][0]: - return [(item['question'], item['answer']) for item in data['data']] - conversations = [] - if 'content' in data: - content = data['content'] - if isinstance(content, list): - for item in content: - if isinstance(item, dict): - if 'q' in item and 'a' in item: - conversations.append((str(item['q']), str(item['a']))) - elif 'question' in item and 'answer' in item: - conversations.append((str(item['question']), str(item['answer']))) - return conversations - return [] - except json.JSONDecodeError: - print(f"错误：{path} 不是有效的JSON文件") - return [] - except Exception as e: - print(f"加载数据失败（{path}）：{str(e)}") - continue - - print(f"警告：未找到数据文件 {data_path}") - return [] - - -# ==================== 15. 主函数 ==================== -def main(): - print("=" * 50) - print("LLM对话AI训练程序 - 改进版") - print("=" * 50) - - device = 'cuda' if torch.cuda.is_available() else 'cpu' - print(f"使用设备: {device}") - - # 选择训练数据文件 - print("\n可用的训练数据文件（在当前目录）：") - import glob - json_files = glob.glob("*.json") - if json_files: - for i, f in enumerate(json_files, 1): - print(f" {i}. {f}") - else: - print(" 未找到JSON数据文件") - - data_path = input("\n请输入训练数据文件路径（或直接回车使用默认: training_data.json）: ").strip() - data_path = data_path.strip('"').strip("'") - if not data_path: - data_path = 'training_data.json' - - config = { - 'data_path': data_path, - 'model_path': 'llm_model.pth', - 'batch_size': 32 if torch.cuda.is_available() else 4, - 'gradient_accumulation_steps': 2, - 'epochs': 100, - 'd_model': 512, - 'n_heads': 8, - 'n_layers': 8, - 'd_ff': 2048, - 'max_length': 128, - 'dropout': 0.1, - 'use_amp': torch.cuda.is_available(), - 'generation_config': { - 'top_k': 50, - 'top_p': 0.9, - 'temperature': 0.8 - } - } - - if torch.cuda.is_available(): - gpu_name = torch.cuda.get_device_name(0) - gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9 - print(f"🔥 检测到CUDA设备: {gpu_name}") - print(f" 💾 显存总量: {gpu_mem:.2f}GB") - print(f" 📦 批次大小: {config['batch_size']}") - print(f" ⚡ 混合精度训练: {'启用' if config['use_amp'] else '禁用'}") - else: - print("⚠ 未检测到CUDA设备，将使用CPU训练") - - # 加载数据 - print(f"\n加载训练数据: {config['data_path']}") - conversations = load_training_data(config['data_path']) - if not conversations: - print(f"\n错误：无法加载训练数据 {config['data_path']}") - return - print(f"加载了 {len(conversations)} 条对话数据") - - # 构建词汇表 - print("\n构建词汇表...") - vocab = build_vocabulary(conversations) - print(f"词汇表大小: {len(vocab)}") - - # 创建数据集 - dataset = ConversationDataset(conversations, vocab, config['max_length']) - dataloader = DataLoader( - dataset, - batch_size=config['batch_size'], - shuffle=True, - num_workers=4 if torch.cuda.is_available() else 0, - pin_memory=True if torch.cuda.is_available() else False, - persistent_workers=True if torch.cuda.is_available() else False - ) - print(f"数据集大小: {len(dataset)}") - print(f"DataLoader批次数量: {len(dataloader)}") - - # 询问是否继续训练 - print("\n" + "=" * 50) - print("训练模式选择") - print("=" * 50) - - pth_files = glob.glob("*.pth") - if pth_files: - for i, f in enumerate(pth_files, 1): - print(f" {i}. {f}") - else: - print(" 未找到模型文件") - - resume_choice = input("\n是否从已有模型继续训练？(y/n，默认n): ").strip().lower() - resume_training = resume_choice == 'y' or resume_choice == 'yes' - - start_epoch = 1 - resume_model_path = None - if resume_training and pth_files: - print("\n可用的模型文件：") - for i, f in enumerate(pth_files, 1): - print(f" {i}. {f}") - - model_choice = input(f"\n请输入要加载的模型编号（1-{len(pth_files)}）: ").strip() - try: - model_idx = int(model_choice) - 1 - if 0 <= model_idx < len(pth_files): - resume_model_path = pth_files[model_idx] - print(f"\n选择加载模型: {resume_model_path}") - except (ValueError, IndexError): - print("\n无效的选择，将重新开始训练") - resume_training = False - - # 训练前参数配置 - print("\n" + "=" * 50) - print("训练参数配置") - print("=" * 50) - - memory_choice = input("是否启用记忆功能？(Y/n，默认Y): ").strip().lower() - use_memory = memory_choice != 'n' and memory_choice != 'no' - - thinking_choice = input("是否启用深度思考？(Y/n，默认Y): ").strip().lower() - use_thinking = thinking_choice != 'n' and thinking_choice != 'no' - - moe_choice = input("是否启用MoE？(Y/n，默认Y): ").strip().lower() - use_moe = moe_choice != 'n' and moe_choice != 'no' - - if use_thinking: - try: - n_thinking_steps = int(input("思考步骤数 (默认3): ").strip() or "3") - except ValueError: - n_thinking_steps = 3 - else: - n_thinking_steps = 1 - - print(f"\n当前配置:") - print(f"记忆功能: {'已启用' if use_memory else '未启用'}") - print(f"深度思考: {'已启用' if use_thinking else '未启用'}") - print(f"MoE: {'已启用' if use_moe else '未启用'}") - print(f"思考步骤数: {n_thinking_steps}") - - confirm = input("\n是否使用以上配置开始训练？(Y/n，默认Y): ").strip().lower() - if confirm == 'n' or confirm == 'no': - print("训练已取消") - return - - # 创建模型 - print("\n创建LLM模型...") - model = LLMModel( - vocab_size=len(vocab), - d_model=config['d_model'], - n_heads=config['n_heads'], - n_layers=config['n_layers'], - d_ff=config['d_ff'], - max_len=config['max_length'], - dropout=config['dropout'], - use_memory=use_memory, - memory_size=10, - use_thinking=use_thinking, - n_thinking_steps=n_thinking_steps, - use_moe=use_moe, - num_experts=4 if use_moe else 1, - top_k=2 if use_moe else 1 - ) - print(f"模型配置: d_model={config['d_model']}, n_layers={config['n_layers']}, n_heads={config['n_heads']}") - print(f"记忆功能: {'已启用' if model.use_memory else '未启用'}") - print(f"深度思考: {'已启用' if model.use_thinking else '未启用'}") - print(f"MoE: {'已启用' if model.use_moe else '未启用'}") - - # 创建训练器 - trainer = LLMTrainer(model, device, use_amp=config['use_amp']) - print(f"\n🎮 训练器配置:") - print(f" - 混合精度: {'✓ 启用 (FP16)' if config['use_amp'] else '✗ 禁用 (FP32)'}") - print(f" - 梯度累积: {config['gradient_accumulation_steps']} 步") - - # 加载已有模型 - if resume_training and resume_model_path: - try: - print(f"\n加载模型参数: {resume_model_path}") - checkpoint = torch.load(resume_model_path, map_location=device) - - if 'epoch' in checkpoint: - start_epoch = checkpoint['epoch'] + 1 - print(f"从第 {checkpoint['epoch']} 轮继续") - else: - import re - epoch_match = re.search(r'epoch[_]?(\d+)', resume_model_path) - if epoch_match: - start_epoch = int(epoch_match.group(1)) + 1 - print(f"从文件名推断：从第 {start_epoch} 轮继续训练") - else: - start_epoch = 1 - - trainer.load_model(resume_model_path) - print(f"✓ 模型参数加载成功") - - if start_epoch > 1: - config['epochs'] = config['epochs'] + start_epoch - 1 - print(f"训练轮数调整为: {config['epochs']} 轮") - except Exception as e: - print(f"✗ 加载模型失败: {e}") - resume_training = False - start_epoch = 1 - - # 训练 - print("\n" + "="*60) - print("🚀 开始训练！") - print("="*60) - start_time = time.time() - - try: - for epoch in range(start_epoch, config['epochs'] + 1): - epoch_start = time.time() - print(f"\n{'='*60}") - print(f"📊 Epoch {epoch}/{config['epochs']} 开始") - print(f"{'='*60}") - - avg_loss = trainer.train_epoch(dataloader, epoch, vocab, config) - epoch_time = time.time() - epoch_start - - print(f"\n{'='*60}") - print(f"📈 Epoch {epoch} 完成 | 耗时: {epoch_time:.1f}秒") - print(f' 📉 Loss: {avg_loss:.4f}') - if torch.cuda.is_available(): - print(f" 💾 显存峰值: {torch.cuda.max_memory_allocated()/1e9:.2f}GB") - torch.cuda.reset_peak_memory_stats() - print(f"{'='*60}\n") - - # 每10个epoch保存一次 - if epoch % 10 == 0: - save_path = f'llm_model_epoch_{epoch}.pth' - print(f"💾 保存检查点: {save_path}") - trainer.save_model(save_path, vocab, epoch) - - # 训练完成 - training_time = time.time() - start_time - print("\n" + "="*60) - print("🎉 训练完成！") - print("="*60) - print(f"⏱️ 总耗时: {training_time:.1f}秒 ({training_time/60:.1f}分钟)") - print("="*60) - - # 保存最终模型 - print(f"💾 保存最终模型: {config['model_path']}") - trainer.save_model(config['model_path'], vocab, config['epochs']) - - except KeyboardInterrupt: - print("\n检测到中断信号，正在保存模型...") - timestamp = time.strftime("%Y%m%d_%H%M%S") - trainer.save_model(f'llm_model_interrupted_{timestamp}.pth', vocab) - print("模型已保存，训练终止") - - # 测试 - print("\n" + "=" * 50) - print("测试模型") - print("=" * 50) - - llm_chat = LLMChat(model, vocab, device, config['max_length']) - - test_questions = ["你好", "你是谁", "谢谢", "再见"] - for q in test_questions: - response = llm_chat.generate_response(q) - print(f"\n问题: {q}") - print(f"回答: {response}") - - # 交互模式 - print("\n" + "=" * 50) - print("进入交互对话模式") - print("=" * 50) - print("可用命令：") - print(" quit or 退出 - 退出对话") - print(" clear or 清除 - 清除对话历史") - print("=" * 50) - - while True: - user_input = input("\n你: ").strip() - if user_input.lower() in ['quit', '退出', '再见']: - print("再见！") - break - if user_input.lower() in ['clear', '清除']: - llm_chat.clear_history() - print("对话历史已清除") - continue - if not user_input: - continue - - response = llm_chat.generate_response(user_input) - print(f"AI: {response}") - - -if __name__ == "__main__": - main() diff --git a/scripts/api_correction.py b/scripts/api_correction.py deleted file mode 100644 index 5e6c4b0..0000000 --- a/scripts/api_correction.py +++ /dev/null @@ -1,148 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -API生成模块 — Astron 5并发（同步根目录版本） -""" -import json, requests, time, os -from concurrent.futures import ThreadPoolExecutor, as_completed - -ASTRON_URL = "https://maas-coding-api.cn-huabei-1.xf-yun.com/v2/chat/completions" -ASTRON_KEY = os.environ.get("ASTRON_KEY", "") -ASTRON_MODEL = "astron-code-latest" - -SYSTEM_PROMPT = """你是一个名叫小狸的8岁中国小女孩。 - -你的人设： -- 8岁女孩子，有点害羞但很可爱 -- 说话简单自然，符合小朋友的语气 -- 回答问题要简短，1-2句话就够了 -- 不要用"喵"结尾，不要过度卖萌 - -回答要求： -1. 语气自然，像真实的小朋友说话 -2. 不懂就说"不知道" -3. 回答简洁但要说清楚，不要讲一半就停""" - - -def call_astron(q, temperature=0.85): - for attempt in range(3): - try: - headers = {"Content-Type": "application/json", "Authorization": f"Bearer {ASTRON_KEY}"} - payload = { - "model": ASTRON_MODEL, - "messages": [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": q} - ], - "temperature": temperature, - "max_tokens": 8192, - } - resp = requests.post(ASTRON_URL, headers=headers, json=payload, timeout=60) - if resp.status_code == 200: - content = resp.json()['choices'][0]['message']['content'].strip() - if content: - return content - else: - time.sleep(1) - except Exception: - time.sleep(2) - return None - - -def batch_generate(questions, temperature=0.85, max_workers=5): - total = len(questions) - results = [] - def call(q): - ans = call_astron(q, temperature) - return (q, ans) if ans else None - with ThreadPoolExecutor(max_workers=max_workers) as executor: - futures = {executor.submit(call, item): i for i, item in enumerate(questions)} - done = 0 - t0 = time.time() - for f in as_completed(futures): - done += 1 - r = f.result() - if r: - results.append({"question": r[0], "answer": r[1]}) - if done % 20 == 0 or done == total: - elapsed = max(time.time() - t0, 1) - rate = len(results) / elapsed * 60 - print(f" Astron: {done}/{total} (已获{len(results)}条, {rate:.0f}条/分)", end="\r") - elapsed = time.time() - t0 - rate = len(results) / max(elapsed, 1) * 60 - print(f" Astron: {total}/{total} → 共{len(results)}条 ({elapsed:.0f}s, {rate:.0f}条/分)") - return results - - -batch_generate_astron = batch_generate -batch_generate_mimo = batch_generate -batch_generate_direct = batch_generate -batch_generate_dual = batch_generate - - -def batch_convert(items, progress=None, max_workers=5): - if not items: - return [] - total = len(items) - results = [] - def convert(item): - q, model_ans = item - prompt = f"问题：{q}\n模型回答：{model_ans}\n\n请用8岁小狸的语气，重新生成更好的回答。保持简短自然，1-2句话。" - ans = call_astron(prompt, temperature=0.3) - return {"question": q, "answer": ans} if ans and len(ans) >= 3 else None - with ThreadPoolExecutor(max_workers=max_workers) as executor: - futures = {executor.submit(convert, item): i for i, item in enumerate(items)} - done = 0 - t0 = time.time() - for f in as_completed(futures): - done += 1 - r = f.result() - if r: - results.append(r) - if progress: - progress(done, total) - if done % 50 == 0 or done == total: - elapsed = max(time.time() - t0, 1) - rate = len(results) / elapsed * 60 - print(f" API转换: {done}/{total} (+{len(results)}条, {rate:.0f}条/分)", end="\r") - elapsed = time.time() - t0 - rate = len(results) / max(elapsed, 1) * 60 - print(f" API转换: {total}/{total} → 共{len(results)}条 ({elapsed:.0f}s, {rate:.0f}条/分)") - return results - - -def batch_generate_multi_temp(questions, temperatures=None, max_workers=5): - if temperatures is None: - temperatures = [0.7, 0.85, 1.0] - if not questions: - return [] - total = len(questions) * len(temperatures) - results = [] - def call_with_temp(q, temp): - ans = call_astron(q, temperature=temp) - return {"question": q, "answer": ans, "temperature": temp} if ans else None - tasks = [(q, t) for q in questions for t in temperatures] - with ThreadPoolExecutor(max_workers=max_workers) as executor: - futures = {executor.submit(call_with_temp, q, t): i for i, (q, t) in enumerate(tasks)} - done = 0 - t0 = time.time() - for f in as_completed(futures): - done += 1 - r = f.result() - if r: - results.append(r) - if done % 50 == 0 or done == total: - elapsed = max(time.time() - t0, 1) - rate = len(results) / elapsed * 60 - print(f" 多温生成: {done}/{total} (+{len(results)}条, {rate:.0f}条/分)", end="\r") - seen = set() - unique = [] - for item in results: - key = (item["question"], item["answer"]) - if key not in seen: - seen.add(key) - unique.append({"question": item["question"], "answer": item["answer"]}) - elapsed = time.time() - t0 - rate = len(unique) / max(elapsed, 1) * 60 - print(f" 多温生成: 共{len(unique)}条 (去重后, {elapsed:.0f}s, {rate:.0f}条/分)") - return unique \ No newline at end of file diff --git a/scripts/daily_data_gen.py b/scripts/daily_data_gen.py deleted file mode 100644 index 83bb626..0000000 --- a/scripts/daily_data_gen.py +++ /dev/null @@ -1,221 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -每日数据生成 — 每天12:00跑 -用 Astron 生成大量小狸风格问答数据，追加到训练集（不触发训练） -""" -import sys, os, json, time, random - -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1) - -from api_correction import batch_generate - -MODEL_DIR = os.path.dirname(os.path.abspath(__file__)) -MAX_WORKERS = 5 # Astron 5并发（无限量） -BATCH_SIZE = 500 - -# 每日生成量 — Astron无限量，直接拉满！ -DAILY_TOPIC_QS = 2000 # 话题衍生 -DAILY_NATURAL_QS = 1000 # 自然闲聊 - -# ===== 种子话题（每天随机选一批让小狸聊） ===== -SEED_TOPICS = [ - "幼儿园", "学校", "朋友", "玩具", "动画片", "宠物", "零食", "天气", - "节日", "生日", "爸爸妈妈", "老师", "作业", "画画", "唱歌", "跳舞", - "游乐园", "动物园", "海洋馆", "公园", "滑滑梯", "秋千", "积木", "拼图", - "自行车", "游泳", "跳绳", "踢球", "冰淇淋", "蛋糕", "水果", "蔬菜", - "彩虹", "星星", "月亮", "太阳", "花", "草", "树", "小动物", - "蝴蝶", "小鸟", "小鱼", "乌龟", "兔子", "猫", "狗", "恐龙", - "过年", "压岁钱", "红包", "新衣服", "放鞭炮", "灯笼", "元宵", - "圣诞老人", "礼物", "糖果", "巧克力", "饼干", "牛奶", "早餐", - "为什么天空是蓝的", "为什么树叶会掉", "为什么有影子", "为什么会下雨", - "为什么会有白天黑夜", "为什么海水是咸的", "为什么会打雷", - "1+1等于几", "数数", "认字", "背诗", "唱歌", "讲故事", - "礼貌", "分享", "排队", "帮助别人", "说谢谢", "对不起", - "害怕", "开心", "难过", "生气", "勇敢", "害羞", - # 拓展话题 - "奥特曼", "小猪佩奇", "汪汪队", "海底小纵队", "超级飞侠", - "白雪公主", "灰姑娘", "小红帽", "三只小猪", "匹诺曹", - "茉莉花", "向日葵", "蒲公英", "含羞草", "仙人掌", - "企鹅", "熊猫", "长颈鹿", "大象", "海豚", "松鼠", "刺猬", - "草莓", "西瓜", "葡萄", "芒果", "蓝莓", "橙子", - "汉堡", "薯条", "披萨", "面条", "饺子", "馄饨", "粽子", - "沙滩", "大海", "高山", "森林", "瀑布", "溶洞", - "火箭", "飞机", "高铁", "轮船", "自行车", "滑板", - "围棋", "象棋", "跳棋", "五子棋", "扑克牌", - "机器人", "变形金刚", "乐高", "磁力片", "橡皮泥", - "春天", "夏天", "秋天", "冬天", "刮风", "下雪", "打雷", - "刷牙", "洗脸", "洗澡", "剪指甲", "梳头", "穿衣服", - "诚实", "分享", "礼貌", "排队", "帮忙", "道歉", "谢谢", - "高兴", "生气", "难过", "害怕", "惊讶", "害羞", "委屈", - "萤火虫", "蜻蜓", "蜜蜂", "蝉", "蜗牛", "蚂蚁", "蜘蛛", - "玉米", "土豆", "番茄", "胡萝卜", "白菜", "茄子", "辣椒", - "毛笔", "铅笔", "蜡笔", "水彩笔", "画板", "颜料", - "竹笋", "蘑菇", "松树", "枫叶", "荷花", "梅花", - "元旦", "春节", "元宵节", "清明节", "端午节", "中秋节", - "万圣节", "感恩节", "圣诞节", "儿童节", "母亲节", "父亲节", - "月球", "太阳系", "银河", "彗星", "流星", "黑洞", - "恐龙", "猛犸象", "剑齿虎", "翼龙", "三角龙", "霸王龙", - "三字经", "弟子规", "唐诗", "宋词", "成语", - "捉迷藏", "老鹰捉小鸡", "丢手绢", "跳房子", "踢毽子", - "轮滑", "滑冰", "滑雪", "冲浪", "攀岩", - "钢琴", "吉他", "小提琴", "笛子", "鼓", "古筝", - "警察", "消防员", "医生", "护士", "科学家", "宇航员", - "红绿灯", "斑马线", "安全带", "头盔", "救生圈", - "原来如此", "没关系", "太好了", "加油", "真棒", "别担心", - "手影", "折纸", "翻花绳", "吹泡泡", "放风筝", - "长城", "故宫", "天安门", "兵马俑", "大熊猫基地", - "小木屋", "城堡", "树屋", "帐篷", "房车", -] - -def generate_questions_from_topics(topics, count=500): - """用话题种子生成自然提问""" - questions = [] - for topic in topics[:count]: - patterns = [ - f"小狸，你喜欢{topic}吗？", - f"给我讲讲{topic}吧", - f"{topic}是什么呀？", - f"你最喜欢什么样的{topic}？", - f"关于{topic}你知道什么？", - f"小狸，{topic}好玩吗？", - f"能不能教我关于{topic}的事情？", - f"你上次玩{topic}是什么时候？", - f"大家为什么都喜欢{topic}呀？", - f"小狸，你在{topic}上有什么有趣的事吗？", - ] - q = random.choice(patterns) - questions.append(q) - return questions - -def generate_natural_questions(count=200): - """生成自然闲聊问题（固定一批高质量问题）""" - base_questions = [ - "小狸你好呀", "今天开心吗", "你几岁了", "你叫什么名字", - "你吃饭了吗", "我们一起玩吧", "你在干嘛呢", "你今天做了什么", - "你喜欢什么颜色", "你最好的朋友是谁", "你最喜欢什么动物", - "你最喜欢吃什么", "你最喜欢什么水果", "你最爱看什么动画片", - "可以给我讲个故事吗", "唱首歌给我听", "教我背一首诗吧", - "你会画画吗", "我们来玩猜谜语", "你家里有谁呀", - "你的房间是什么样子的", "你书包里有什么", "你最喜欢上什么课", - "今天天气真好", "外面下雨了", "好冷啊", "好热啊", - "我有点不开心", "我今天好开心", "我害怕打雷", - "你睡午觉了吗", "晚安", "早安", "快来", - "我想吃冰淇淋", "我想去游乐场", "我想养一只小猫", - "为什么雪是白色的", "为什么天是蓝的", "为什么花是香的", - "月亮上有人吗", "星星为什么会眨眼睛", "大海有多大", - "你会数数吗", "1+1等于几", "100+200等于多少", - "我喜欢你小狸", "你好可爱", "我们做朋友吧", - "你今天真好看", "谢谢你", "你真好", - "对不起", "我不小心摔倒了", "我被人欺负了", - "怎么办呀", "帮帮我", "我该怎么做", - "什么是爱", "什么是勇敢", "什么是善良", - "长大是什么感觉", "为什么要上学", "为什么要听话", - "你会一直陪着我吗", "你想去什么地方玩", - ] - return random.sample(base_questions, min(count, len(base_questions))) - -def main(): - print("=" * 60) - print("📅 每日数据生成") - print(f"时间: {time.strftime('%Y-%m-%d %H:%M:%S')}") - print("=" * 60) - - # 1. 混合生成问题 - random.seed(int(time.time())) - n_topic_qs = DAILY_TOPIC_QS # 话题衍生 - n_natural_qs = DAILY_NATURAL_QS # 自然闲聊 - - random.shuffle(SEED_TOPICS) - topic_questions = generate_questions_from_topics(SEED_TOPICS, n_topic_qs) - natural_questions = generate_natural_questions(n_natural_qs) - - all_questions = topic_questions + natural_questions - random.shuffle(all_questions) - - print(f"\n准备生成 {len(all_questions)} 条问答 (话题衍生: {len(topic_questions)}, 自然: {len(natural_questions)})") - - # 2. 用不同温度分批生成，增加多样性 - TEMPS = [0.85, 0.9, 0.95, 1.0] - all_results = [] - existing = set() - - # 加载已有数据用于去重 - train_path = os.path.join(MODEL_DIR, 'xiaoli_train_v2.json') - with open(train_path, encoding='utf-8') as f: - cur_train = json.load(f) - for it in cur_train: - key = (it.get('question','')[:100], it.get('answer','')[:100]) - existing.add(key) - print(f"已有训练集: {len(cur_train)} 条") - - # 分批+变温度 - total_batches = (len(all_questions) + BATCH_SIZE - 1) // BATCH_SIZE - t_start = time.time() - - for batch_idx in range(total_batches): - batch_qs = all_questions[batch_idx * BATCH_SIZE : (batch_idx + 1) * BATCH_SIZE] - temp = random.choice(TEMPS) - - print(f"\n--- 批次 {batch_idx + 1}/{total_batches} ({len(batch_qs)}条, temp={temp}) ---") - - t0 = time.time() - results = batch_generate(batch_qs, temperature=temp, max_workers=MAX_WORKERS) - elapsed = time.time() - t0 - - new_in_batch = 0 - for r in results: - key = (r['question'][:100], r['answer'][:100]) - if key not in existing: - existing.add(key) - all_results.append(r) - new_in_batch += 1 - - total_elapsed = time.time() - t_start - rate = len(all_results) / max(total_elapsed, 1) * 60 - print(f" 本批新增 {new_in_batch} 条 | 累计 {len(all_results)} 条 | {rate:.0f}条/分") - - if not all_results: - print("\n⚠️ 没有生成新数据") - return - - # 3. 追加到训练集 - print(f"\n追加 {len(all_results)} 条到训练集...") - for tgt in ['xiaoli_train_v2.json', 'xiaoli_merged_data_v2.json']: - tp = os.path.join(MODEL_DIR, tgt) - with open(tp, encoding='utf-8') as f: - cur = json.load(f) - - ex = set((it.get('question','')[:100], it.get('answer','')[:100]) for it in cur) - to_add = [] - for r in all_results: - key = (r['question'][:100], r['answer'][:100]) - if key not in ex: - ex.add(key) - to_add.append({'question': r['question'], 'answer': r['answer']}) - - cur.extend(to_add) - with open(tp, 'w', encoding='utf-8') as f: - json.dump(cur, f, ensure_ascii=False, indent=2) - - if 'train' in tgt: - print(f" ✅ 训练集 +{len(to_add)} 条 → {len(cur)} 条") - else: - print(f" ✅ 合并集 +{len(to_add)} 条 → {len(cur)} 条") - - # 4. 保存今日生成日志 - today = time.strftime('%Y-%m-%d') - log_path = os.path.join(MODEL_DIR, f'daily_gen_{today}.json') - with open(log_path, 'w', encoding='utf-8') as f: - json.dump(all_results, f, ensure_ascii=False, indent=2) - - elapsed_total = time.time() - t_start - print(f"\n{'='*60}") - print(f"✅ 今日生成完成! {len(all_results)} 条新数据") - print(f"用时: {elapsed_total/60:.1f} 分钟") - print(f"日志: daily_gen_{today}.json") - print(f"{'='*60}") - -if __name__ == '__main__': - main() \ No newline at end of file diff --git a/scripts/generate_batch.py b/scripts/generate_batch.py deleted file mode 100644 index 3ef9cce..0000000 --- a/scripts/generate_batch.py +++ /dev/null @@ -1,57 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -单批生成 — 从mega_questions.json取指定batch生成并追加 -用法: python3 generate_batch.py -""" -import sys, os, json, time -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1) - -from api_correction import batch_generate - -model_dir = os.path.dirname(os.path.abspath(__file__)) -batch_num = int(sys.argv[1]) -batch_size = int(sys.argv[2]) - -# 加载问题列表 -with open(os.path.join(model_dir, 'mega_questions.json'), 'r', encoding='utf-8') as f: - mega = json.load(f) - -# 计算本批范围 -start = (batch_num - 1) * batch_size -end = min(start + batch_size, len(mega)) -chunk = [item['question'] for item in mega[start:end]] - -if not chunk: - print(f"批{batch_num}: 无数据") - sys.exit(0) - -print(f"批 {batch_num}: {len(chunk)}题 (序号{start}-{end})") -t0 = time.time() - -results = batch_generate(chunk, temperature=0.85) -dt = time.time() - t0 - -# 保存中间文件 -interim_path = os.path.join(model_dir, f'mega_interim_{batch_num}.json') -with open(interim_path, 'w', encoding='utf-8') as f: - json.dump(results, f, ensure_ascii=False, indent=2) - -# 追加到训练集 -for tgt in ['xiaoli_train_v2.json', 'xiaoli_merged_data_v2.json']: - tgt_path = os.path.join(model_dir, tgt) - if not os.path.exists(tgt_path): - continue - with open(tgt_path, 'r', encoding='utf-8') as f: - cur = json.load(f) - existing = set((item.get("question",""), item.get("answer","")) for item in cur) - truly_new = [item for item in results if (item["question"], item["answer"]) not in existing] - if truly_new: - cur.extend(truly_new) - with open(tgt_path, 'w', encoding='utf-8') as f: - json.dump(cur, f, ensure_ascii=False, indent=2) - if tgt == 'xiaoli_train_v2.json': - print(f" ✅ 追加训练集: +{len(truly_new)}条, 现有{len(cur)}条") - -print(f"批{batch_num}完成: +{len(results)}条 ({dt:.0f}s)") \ No newline at end of file diff --git a/scripts/generate_diverse.py b/scripts/generate_diverse.py deleted file mode 100644 index a784b6d..0000000 --- a/scripts/generate_diverse.py +++ /dev/null @@ -1,73 +0,0 @@ -#!/usr/bin/env python3 -"""高温度多样化生成 — 每个问题跑 N 次，获得多种回答""" -import sys, os, json, random, time -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1) -from api_correction import batch_generate_astron as batch_generate - -MODEL_DIR = os.path.dirname(os.path.abspath(__file__)) - -# 载入全部问题 -with open(os.path.join(MODEL_DIR, 'mega_questions.json')) as f: - mega = json.load(f) -total_q = len(mega) -print(f"mega_questions 共 {total_q} 条") - -# 温度配置：逐渐升高，产出多样化 -TEMP_CONFIG = [0.9, 1.1, 1.3, 1.5] # 跑4轮 -all_new_results = [] - -for round_idx, temp in enumerate(TEMP_CONFIG): - print(f"\n{'='*60}") - print(f"第{round_idx+1}轮: 温度={temp}") - print(f"{'='*60}") - - questions = [item['question'] for item in mega] - - t0 = time.time() - results = batch_generate(questions, temperature=temp) - elapsed = time.time() - t0 - - # 标注来源 - for r in results: - r['temperature'] = temp - r['round'] = round_idx + 1 - - all_new_results.extend(results) - print(f"轮{round_idx+1} 完成: {len(results)}条 ({elapsed:.0f}s)") - print(f"累计: {len(all_new_results)}条") - -# 去重（问题+回答完全重复才去） -print(f"\n总体去重中...") -seen = set() -unique_results = [] -for r in all_new_results: - key = (r['question'][:100], r['answer'][:100]) - if key not in seen: - seen.add(key) - unique_results.append(r) - -print(f"去重前: {len(all_new_results)}条 → 去重后: {len(unique_results)}条") - -# 追加到训练集 -for tgt in ['xiaoli_train_v2.json', 'xiaoli_merged_data_v2.json']: - tp = os.path.join(MODEL_DIR, tgt) - with open(tp, encoding='utf-8') as f: - cur = json.load(f) - - ex = set((it.get("question","")[:100], it.get("answer","")[:100]) for it in cur) - to_add = [] - for r in unique_results: - key = (r['question'][:100], r['answer'][:100]) - if key not in ex: - ex.add(key) - to_add.append({"question": r['question'], "answer": r['answer']}) - - cur.extend(to_add) - with open(tp, 'w', encoding='utf-8') as f: - json.dump(cur, f, ensure_ascii=False, indent=2) - if 'train' in tgt: - print(f"✅ 训练集 +{len(to_add)}条 → {len(cur)}条") - -print(f"\n🎉 全部完成!") -print(f"总数: {len(cur)}条") \ No newline at end of file diff --git a/scripts/generate_diverse_single.py b/scripts/generate_diverse_single.py deleted file mode 100644 index 2218106..0000000 --- a/scripts/generate_diverse_single.py +++ /dev/null @@ -1,132 +0,0 @@ -#!/usr/bin/env python3 -""" -高温度多样化生成 — 单线程，与常规跑批共用 Astron 5并发中的1个 -对已完成批1-6的问题，用高温度重新生成，获得多样回答 -""" -import sys, os, json, random, time, requests -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1) - -ASTRON_URL = "https://maas-coding-api.cn-huabei-1.xf-yun.com/v2/chat/completions" -ASTRON_KEY = os.environ.get("ASTRON_KEY", "") -ASTRON_MODEL = "astron-code-latest" -MODEL_DIR = os.path.dirname(os.path.abspath(__file__)) - -SYSTEM_PROMPT = """你是一个名叫小狸的8岁中国小女孩。 - -你的人设： -- 8岁女孩子，有点害羞但很可爱 -- 说话简单自然，符合小朋友的语气 -- 回答问题要简短，1-2句话就够了 -- 不要用"喵"结尾，不要过度卖萌 - -回答要求： -1. 语气自然，像真实的小朋友说话 -2. 不懂就说"不知道" -3. 回答简洁但要说清楚，不要讲一半就停""" - - -def call_astron(q, temperature): - for attempt in range(3): - try: - headers = {"Content-Type": "application/json", "Authorization": f"Bearer {ASTRON_KEY}"} - payload = { - "model": ASTRON_MODEL, - "messages": [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": q} - ], - "temperature": temperature, - "max_tokens": 8192, - } - resp = requests.post(ASTRON_URL, headers=headers, json=payload, timeout=60) - if resp.status_code == 200: - content = resp.json()['choices'][0]['message']['content'].strip() - if content: - return content - else: - time.sleep(1) - except Exception: - time.sleep(2) - return None - - -def main(): - # 从已完成批载入问题 - all_questions = [] - for i in range(1, 7): - fp = os.path.join(MODEL_DIR, f'mega_interim_{i}.json') - if os.path.exists(fp): - with open(fp) as f: - items = json.load(f) - for item in items: - all_questions.append(item['question']) - print(f"批{i}: {len(items)}条已载入") - - # 去重 - all_questions = list(dict.fromkeys(all_questions)) - print(f"共 {len(all_questions)} 个去重问题") - - # 温度配置 - temp_configs = [0.9, 1.3, 1.5] - all_new = [] - - for round_idx, temp in enumerate(temp_configs): - print(f"\n{'='*50}") - print(f"轮{round_idx+1}/3: 温度={temp}") - print(f"{'='*50}") - - # 单线程顺序跑 - results = [] - total = len(all_questions) - t0 = time.time() - - for idx, q in enumerate(all_questions): - ans = call_astron(q, temp) - if ans: - results.append({"question": q, "answer": ans}) - - if (idx + 1) % 20 == 0 or idx + 1 == total: - elapsed = max(time.time() - t0, 1) - rate = len(results) / elapsed * 60 - eta = (total - idx - 1) / max(idx + 1, 1) * elapsed / 60 - print(f" 多样: {idx+1}/{total} (已获{len(results)}条, {rate:.0f}条/分, ETA:{eta:.0f}分)", end="\r") - - elapsed = time.time() - t0 - rate = len(results) / max(elapsed, 1) * 60 - print(f"\n 轮{round_idx+1}完成: {len(results)}条 ({elapsed:.0f}s, {rate:.0f}条/分)") - - for r in results: - r['temperature'] = temp - r['round'] = round_idx + 1 - all_new.extend(results) - - # ✅ 每轮跑完立即去重并追加到训练集 - seen = set() - unique = [] - for r in all_new: - key = (r['question'][:100], r['answer'][:100]) - if key not in seen: - seen.add(key) - unique.append(r) - - for tgt in ['xiaoli_train_v2.json', 'xiaoli_merged_data_v2.json']: - tp = os.path.join(MODEL_DIR, tgt) - with open(tp, encoding='utf-8') as f: - cur = json.load(f) - ex = set((it.get("question","")[:100], it.get("answer","")[:100]) for it in cur) - to_add = [] - for r in unique: - key = (r['question'][:100], r['answer'][:100]) - if key not in ex: - ex.add(key) - to_add.append({"question": r['question'], "answer": r['answer']}) - cur.extend(to_add) - with open(tp, 'w', encoding='utf-8') as f: - json.dump(cur, f, ensure_ascii=False, indent=2) - if 'train' in tgt: - print(f"✅ 轮{round_idx+1} 训练集 +{len(to_add)}条 → {len(cur)}条") - - -if __name__ == '__main__': - main() \ No newline at end of file diff --git a/scripts/generate_fresh_qa.py b/scripts/generate_fresh_qa.py deleted file mode 100644 index 0693a26..0000000 --- a/scripts/generate_fresh_qa.py +++ /dev/null @@ -1,240 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -全新问答对生成 — Astron 凭空造（新问题+新回答），5并发 -不依赖现有问题列表，从种子话题衍生 -增量保存防丢 -""" -import json, requests, time, os, sys, random, hashlib -from concurrent.futures import ThreadPoolExecutor, as_completed -sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1) - -# ====== Astron 配置 ====== -ASTRON_URL = "https://maas-coding-api.cn-huabei-1.xf-yun.com/v2/chat/completions" -ASTRON_KEY = os.environ.get("ASTRON_KEY", "") -ASTRON_MODEL = "astron-code-latest" - -SYSTEM_PROMPT = """你是一个名叫小狸的8岁中国小女孩。 - -你的人设： -- 8岁女孩子，有点害羞但很可爱 -- 说话简单自然，符合小朋友的语气 -- 回答问题要简短，1-2句话就够了 -- 不要用"喵"结尾，不要过度卖萌 - -回答要求： -1. 语气自然，像真实的小朋友说话 -2. 不懂就说"不知道" -3. 回答简洁但要说清楚""" - -# ====== 种子场景（生成对话的起点） ====== -SCENES = [ - # 日常场景 - ("问候", ["打招呼", "问好", "早上起来", "放学回家", "见到朋友"]), - ("情绪", ["开心", "难过", "生气", "无聊", "害怕", "惊喜", "委屈", "害羞"]), - ("分享", ["今天发生的事", "周末做了什么", "假期计划", "新玩具", "考试成绩"]), - ("喜好", ["喜欢的颜色", "喜欢的动物", "喜欢的食物", "喜欢的动画片", "喜欢的游戏"]), - ("好奇", ["为什么问题", "想知道的事", "不懂的事", "想象的事", "奇怪的问题"]), - ("日常", ["起床", "刷牙", "吃饭", "上学", "做作业", "睡觉"]), - ("朋友", ["好朋友", "吵架", "一起玩", "分享秘密", "帮忙"]), - ("家庭", ["爸爸妈妈", "爷爷奶奶", "兄弟姐妹", "家人一起", "家庭活动"]), - ("想象", ["童话故事", "梦想", "如果我有魔法", "长大想做什么", "幻想朋友"]), - ("安慰", ["受伤了", "丢东西", "被骂了", "考试没考好", "被欺负", "做噩梦"]), - ("节日", ["生日", "过年", "圣诞节", "儿童节", "中秋"]), - ("自然", ["天气", "动物", "植物", "星空", "大海", "四季"]), - ("学校", ["老师", "同学", "上课", "课间", "课外活动", "考试"]), - ("晚安", ["说晚安", "睡前", "梦", "明天见", "晚安故事"]), -] - -def call_astron(prompt, temperature=0.9): - for attempt in range(3): - try: - headers = {"Content-Type": "application/json", "Authorization": f"Bearer {ASTRON_KEY}"} - payload = { - "model": ASTRON_MODEL, - "messages": [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": prompt} - ], - "temperature": temperature, - "max_tokens": 4096, - } - resp = requests.post(ASTRON_URL, headers=headers, json=payload, timeout=60) - if resp.status_code == 200: - content = resp.json()['choices'][0]['message']['content'].strip() - if content and len(content) > 10: - return content - else: - time.sleep(1) - except Exception: - time.sleep(2) - return None - - -def parse_qa_pairs(text): - """解析 Astron 返回的问答对列表""" - lines = text.strip().split('\n') - pairs = [] - current_q = None - for line in lines: - line = line.strip() - if not line: - continue - q_match = __import__('re').match(r'^[Q问][：:]\s*(.*)', line) - a_match = __import__('re').match(r'^[A答][：:]\s*(.*)', line) - if q_match: - current_q = q_match.group(1).strip() - elif a_match and current_q: - a_text = a_match.group(1).strip() - if current_q and a_text and len(a_text) > 2 and len(current_q) > 2: - pairs.append((current_q, a_text)) - current_q = None - return pairs - - -def pair_fingerprint(q, a): - """QA对指纹""" - raw = f"{q[:40]}|{a[:40]}" - return hashlib.md5(raw.encode()).hexdigest() - - -def generate_batch_pairs(batch_size=10, existing_fps=None): - """生成一批全新的问答对""" - if existing_fps is None: - existing_fps = set() - - # 随机选场景 - category, sub_scenes = random.choice(SCENES) - sub = random.choice(sub_scenes) - # 随机生成要求 - style = random.choice(["日常聊天", "撒娇可爱", "有点害羞", "好奇宝宝", "暖心安慰"]) - mood = random.choice(["开心", "平静", "有点小情绪", "好奇", "温暖"]) - - prompt = f"""请以小狸的身份，生成{batch_size}个不同的日常对话问答对。 - -对话背景： -- 场景：{category}（{sub}） -- 风格：{style} -- 小狸心情：{mood} - -要求： -1. 每个问答对包含一问一答，问是用户说的，答是小狸回答 -2. 问要自然，像真实用户会说的话（不要提"小狸"） -3. 答要是小狸的风格，1-2句话，简短自然 -4. {batch_size}个问答要对同一个话题但角度不同 -5. 不要重复一模一样的问法 - -请按以下格式输出： -Q：用户提问内容 -A：小狸回答内容 ----（分隔） -Q：用户提问内容 -A：小狸回答内容 ----""" - - result = call_astron(prompt, temperature=random.uniform(0.85, 1.0)) - if not result: - return [] - - pairs = parse_qa_pairs(result) - - # 去重 - deduped = [] - for q, a in pairs: - fp = pair_fingerprint(q, a) - if fp not in existing_fps: - existing_fps.add(fp) - deduped.append({'question': q, 'answer': a}) - - return deduped - - -def main(): - output_dir = '/workspace/xiaoli-llm-chat/fresh_qa_outputs' - os.makedirs(output_dir, exist_ok=True) - - # ====== 恢复进度 ====== - progress_file = os.path.join(output_dir, 'fresh_qa_progress.json') - all_pairs = [] - if os.path.exists(progress_file): - with open(progress_file, encoding='utf-8') as f: - all_pairs = json.load(f) - print(f"已有进度: {len(all_pairs)} 条") - - # 加载已有训练集做去重 - train_path = '/workspace/xiaoli-llm-chat/xiaoli_train_v2.json' - with open(train_path, encoding='utf-8') as f: - train_data = json.load(f) - existing_fps = set(pair_fingerprint(item.get('question',''), item.get('answer','')) for item in train_data) - for p in all_pairs: - existing_fps.add(pair_fingerprint(p['question'], p['answer'])) - print(f"去重指纹池: {len(existing_fps)} 个") - - # ====== 参数 ====== - TARGET = 500000 # 目标 50 万条全新问答 - BATCH_SIZE = 15 # 每批次生成 15 条（更多问题，效率更高） - BATCHES_PER_SAVE = 20 # 每 20 批次保存一次 - MAX_WORKERS = 5 # Astron 5并发 - - remaining = TARGET - len(all_pairs) - total_iterations = (remaining + BATCH_SIZE - 1) // BATCH_SIZE - print(f"目标: {TARGET:,} 条, 还需: {remaining:,} 条") - print(f"轮次: {total_iterations} 轮 (每轮{BATCH_SIZE}条, {MAX_WORKERS}并发)") - print(f"{'='*60}\n") - - t_start = time.time() - batch_counter = 0 - - for iteration in range(total_iterations): - t0 = time.time() - - # 并发生成 - with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor: - futures = [ - executor.submit(generate_batch_pairs, BATCH_SIZE, existing_fps) - for _ in range(MAX_WORKERS) - ] - for future in as_completed(futures): - try: - new_pairs = future.result() - if new_pairs: - for p in new_pairs: - existing_fps.add(pair_fingerprint(p['question'], p['answer'])) - all_pairs.append(p) - except Exception: - pass - - batch_counter += 1 - elapsed = time.time() - t_start - rate = len(all_pairs) / max(elapsed, 1) * 60 - - # 定期保存 - if batch_counter % BATCHES_PER_SAVE == 0 or iteration == total_iterations - 1: - with open(progress_file, 'w', encoding='utf-8') as f: - json.dump(all_pairs, f, ensure_ascii=False, indent=2) - - eta = (TARGET - len(all_pairs)) / max(rate, 0.1) / 60 if rate > 0 else 0 - - if iteration % 5 == 0 or iteration == total_iterations - 1: - print(f"轮 {iteration+1}/{total_iterations} | " - f"累计 {len(all_pairs):,}/{TARGET:,} | " - f"{rate:.0f}条/分 | " - f"已用 {elapsed/60:.1f}分 | " - f"预计剩余 {eta:.1f}分") - - # ====== 完成 ====== - elapsed_total = time.time() - t_start - print(f"\n{'='*60}") - print(f"全新问答生成完成!") - print(f"总数: {len(all_pairs):,} 条") - print(f"用时: {elapsed_total/60:.1f} 分钟") - print(f"速率: {len(all_pairs)/max(elapsed_total,1)*60:.0f} 条/分") - - # 保存最终 - final_path = '/data/user/work/fresh_qa_v1.json' - with open(final_path, 'w', encoding='utf-8') as f: - json.dump(all_pairs, f, ensure_ascii=False, indent=2) - print(f"保存到: {final_path}") - -if __name__ == '__main__': - main() \ No newline at end of file diff --git a/scripts/generate_massive_general.py b/scripts/generate_massive_general.py deleted file mode 100644 index 8595288..0000000 --- a/scripts/generate_massive_general.py +++ /dev/null @@ -1,389 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -海量通用数据蒸馏 v3 — 不设场景标签，纯自然对话，多策略并行 -目标：500万+ 纯 (question, answer) 对 -""" -import json, requests, time, os, sys, random, re -from concurrent.futures import ThreadPoolExecutor, as_completed -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1) - -MODEL_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -OUT_DIR = os.path.join(MODEL_DIR, 'general_outputs') -os.makedirs(OUT_DIR, exist_ok=True) - -# ====== Astron 配置 ====== -ASTRON_URL = "https://maas-coding-api.cn-huabei-1.xf-yun.com/v2/chat/completions" -ASTRON_KEY = os.environ.get("ASTRON_KEY", "") -ASTRON_MODEL = "astron-code-latest" - -# ====== 主题全覆盖（六大类 × 子类） ====== -TOPICS = { - "日常": [ - "起床洗漱", "吃早餐", "上学路上", "上课", "课间休息", - "午餐时间", "放学回家", "做作业", "晚饭", "睡前时光", - "周末活动", "天气变化", "穿衣打扮", "零花钱", "家务劳动", - "看电视", "玩玩具", "去公园", "逛超市", "去亲戚家", - ], - "情绪": [ - "开心高兴", "难过伤心", "生气愤怒", "害怕恐惧", "紧张焦虑", - "孤独寂寞", "委屈冤枉", "失望沮丧", "感动温暖", "兴奋激动", - "无聊烦闷", "尴尬害羞", "后悔内疚", "羡慕嫉妒", "感激感恩", - ], - "社交": [ - "交新朋友", "和好朋友玩", "和同学相处", "和老师交流", "和兄弟姐妹", - "和爸爸妈妈", "和爷爷奶奶", "和邻居", "和陌生人", "和喜欢的人", - "吵架和好", "被欺负", "被误解", "被表扬", "被批评", - "帮助别人", "请求帮助", "分享东西", "邀请做客", "告别分开", - ], - "知识": [ - "动物世界", "植物花草", "天文星空", "天气自然", "人体健康", - "食物营养", "颜色形状", "数字计算", "汉字词语", "科学现象", - "历史故事", "地理知识", "交通工具", "乐器音乐", "绘画手工", - "运动体育", "节日习俗", "各国文化", "发明创造", "宇宙太空", - ], - "想象": [ - "超能力", "魔法世界", "未来科技", "外星人", "恐龙时代", - "童话故事", "变成动物", "飞行梦", "海底世界", "神秘岛屿", - "时间旅行", "超级英雄", "会说话的玩具", "隐形衣", "变身药水", - "机器人朋友", "彩虹桥", "星星糖果", "云朵房子", "梦中冒险", - ], - "成长": [ - "学新本领", "克服困难", "不怕失败", "坚持努力", "自我肯定", - "梦想未来", "职业理想", "好习惯", "改掉坏毛病", "学会分享", - "学会等待", "学会道歉", "学会感恩", "学会勇敢", "学会理解", - "交心谈话", "秘密分享", "烦恼倾诉", "人生道理", "温暖鼓励", - ], -} - -FLAT_TOPICS = [] -for cat, subs in TOPICS.items(): - for sub in subs: - FLAT_TOPICS.append((cat, sub)) - -# ====== 多样化的开场方式 ====== -OPENERS = [ - "小狸", "嗨小狸", "小狸在吗", "嘿嘿", "那个…", "喂喂喂", - "小狸小狸", "我想问你", "你知道吗", "你说", "我突然想到", - "我跟你讲", "小狸你听我说", "问你个事", "我有个问题", - "好无聊啊", "好开心呀", "好烦", "怎么办呀", - "小狸你在干嘛", "我来了", "出来玩", "跟你说个事", -] - -# ====== 策略1: 直接生成纯问答对 ====== -def make_direct_pairs_prompt(count=15): - """无场景标签，直接生成纯问答对""" - # 随机选3-5个主题混搭 - selected = random.sample(FLAT_TOPICS, min(random.randint(3, 5), len(FLAT_TOPICS))) - topics_str = "、".join([f"{c}({s})" for c, s in selected]) - - prompt = f"""请直接生成 {count} 条聊天问答对。 - -格式要求（严格按此格式）： -问：xxxxxxxxx | 答：xxxxxxxxx - -用户的问题要：口语化、自然、像真人（5-25字） -小狸的回答要：8岁小女孩语气，可爱活泼（10-40字） -话题要多样：{topics_str} - -示例： -问：你怕不怕黑呀？ | 答：有一点点…但是抱着小熊玩偶就不怕啦！ -问：我今天跑步摔了一跤 | 答：疼不疼呀？我给你吹吹就不疼了！ - -现在生成 {count} 条：""" - return prompt - -# ====== 策略2: 短对话（3-5轮） ====== -def make_mini_chat_prompt(count=8): - """微型多轮对话，每段3-5轮""" - ops = random.sample(OPENERS, min(count, len(OPENERS))) - - prompt = f"""请生成 {count} 段微型聊天对话，每段3-5轮（一来一回算1轮）。 - -要求： -- 每段一个独立话题，话题要多样（日常、情绪、知识、想象、社交、成长等） -- 用户说话要特别自然、口语化，像真人随口说的 -- 小狸像8岁小女孩，回答要可爱、有画面感 -- 话题要自然延伸，不要生硬转折 - -格式例子： -话题：日常 -用户：今天好冷啊 -小狸：是呀！我穿了厚厚的羽绒服，像只小企鹅一样～ -用户：哈哈哈真的好像 -小狸：（鼓起腮帮子）才不像呢！企鹅走路摇摇摆摆的，我走路可稳啦！ -... - -现在生成 {count} 段，每段用"话题：xxx"开头：""" - return prompt - -# ====== 策略3: 长对话（6-10轮自由展开） ====== -def make_long_chat_prompt(): - """一场长对话，话题自然展开""" - topic = random.choice(FLAT_TOPICS) - opener = random.choice(OPENERS) - - emotions = ["开心", "难过", "兴奋", "沮丧", "好奇", "担心", "轻松", "紧张"] - emotion = random.choice(emotions) - - prompt = f"""请生成一段自然的日常聊天，{topic[0]}类-{topic[1]}话题。 - -用户开场：{opener} -用户当下的心情：{emotion} - -要求： -- 用户很自然地说出开场白，就像真人发微信一样 -- 小狸正常回复，话题慢慢展开，不要硬聊 -- 聊6-10轮（一问一答算一轮） -- 每轮格式：用户：xxx\n小狸：xxx -- 真实感第一，不要刻意 - -请直接输出对话，不要加说明：""" - return prompt - -# ====== 策略4: 单轮问答（高吞吐） ====== -def make_single_turn_prompt(batch_size=20): - """纯单轮问答，批量生成""" - selected = random.sample(FLAT_TOPICS, min(6, len(FLAT_TOPICS))) - - prompt = f"""生成 {batch_size} 条日常问答对。 - -要求： -1. 用户的话：7-25字，口语化，真实感强 -2. 小狸回答：1-2句，8岁小女孩语气，可爱活泼 -3. 覆盖多样化话题：{', '.join([s for _, s in selected[:6]])}（但不仅限于这些） -4. 每行一条，格式：问: ... | 答: ... -5. 不要重复相同的问题 - -示例： -问: 你喜不喜欢下雨天呀？ | 答: 喜欢！因为可以踩水坑，啪嗒啪嗒的超好玩！ -问: 我今天被同学起外号了 | 答: （歪着头）那…你喜欢这个外号吗？不喜欢的话要大声说出来哦！ -问: 你说世界上有外星人吗？ | 答: 我觉得有！说不定他们正在用望远镜看我们呢～ - -生成：""" - return prompt - -# ====== API调用 ====== -def call_astron(prompt, temp=0.82, max_tokens=4096): - for attempt in range(3): - try: - headers = {"Content-Type": "application/json", "Authorization": f"Bearer {ASTRON_KEY}"} - payload = { - "model": ASTRON_MODEL, - "messages": [ - {"role": "system", "content": "你是一个擅长写日常对话的助手。小狸是一个8岁的小女孩，说话可爱自然。"}, - {"role": "user", "content": prompt} - ], - "temperature": temp, - "max_tokens": max_tokens, - } - resp = requests.post(ASTRON_URL, headers=headers, json=payload, timeout=150) - if resp.status_code == 200: - content = resp.json()['choices'][0]['message']['content'].strip() - if content and len(content) > 20: - return content - else: - time.sleep(0.5 + attempt) - except Exception: - time.sleep(1 + attempt) - return None - -# ====== 解析 ====== -def parse_pairs(text): - """从"问：xxx | 答：xxx"格式解析问答对""" - pairs = [] - for line in text.split('\n'): - line = line.strip() - parts = re.split(r'\s*[|｜]\s*', line) - if len(parts) == 2: - q = re.sub(r'^问\s*[：:]\s*', '', parts[0]).strip() - a = re.sub(r'^答\s*[：:]\s*', '', parts[1]).strip() - if q and a and len(q) >= 2 and len(a) >= 2 and len(a) < 300 and len(q) < 200: - pairs.append((q, a)) - return pairs - -def parse_mini_chats(text): - """解析微型多轮对话""" - pairs = [] - segments = re.split(r'话题[：:].*?\n', text) - - for seg in segments: - seg = seg.strip() - lines = seg.split('\n') - q = None - for line in lines: - line = line.strip() - if not line: - continue - q_match = re.match(r'^(用户|我)[：:]\s*(.*)', line) - a_match = re.match(r'^(小狸|答)[：:]\s*(.*)', line) - - if q_match: - q_text = q_match.group(2).strip() - if q_text and len(q_text) > 1: - q = q_text - elif a_match and q: - a_text = a_match.group(2).strip() - if a_text and len(a_text) > 1 and len(a_text) < 300: - pairs.append((q, a_text)) - q = None - return pairs - -def parse_long_chat(text): - """解析长对话""" - pairs = [] - q = None - for line in text.split('\n'): - line = line.strip() - if not line: - continue - q_match = re.match(r'^(用户|我)[：:]\s*(.*)', line) - a_match = re.match(r'^(小狸)[：:]\s*(.*)', line) - - if q_match: - q_text = q_match.group(2).strip() - if q_text and len(q_text) > 1: - q = q_text - elif a_match and q: - a_text = a_match.group(2).strip() - if a_text and len(a_text) > 1 and len(a_text) < 300: - pairs.append((q, a_text)) - q = None - return pairs - -# ====== 生成器（按策略分发） ====== -GENERATORS = [ - # (name, prompt_fn, parse_fn, batch_size_estimate, weight_for_diversity) - ("direct", make_direct_pairs_prompt, parse_pairs, 15, 35), - ("mini_chat", make_mini_chat_prompt, parse_mini_chats, 20, 25), - ("long_chat", make_long_chat_prompt, parse_long_chat, 20, 20), - ("single", make_single_turn_prompt, parse_pairs, 20, 20), -] - -def run_one_generation(): - """运行一次生成任务（随机选一种策略）""" - # 按权重选策略 - weights = [w for _, _, _, _, w in GENERATORS] - gen = random.choices(GENERATORS, weights=weights, k=1)[0] - name, prompt_fn, parse_fn, batch_est, _ = gen - - prompt = prompt_fn() - result = call_astron(prompt, temp=random.uniform(0.75, 0.92), max_tokens=4096 if 'long' not in name else 6144) - if not result: - return None - - pairs = parse_fn(result) - if not pairs: - return None - - return pairs - -def main(): - print("=" * 65) - print(" 海量通用数据蒸馏 v3") - print(" 不设场景标签 | 6大类120+子话题 | 4种策略") - print("=" * 65) - - # 加载已有进度 - all_pairs = [] - existing_keys = set() - progress_file = os.path.join(OUT_DIR, 'all_general_progress.json') - if os.path.exists(progress_file): - existing = json.load(open(progress_file, encoding='utf-8')) - all_pairs = existing - for p in existing: - existing_keys.add((p.get('question',''), p.get('answer',''))) - print(f"已有进度: {len(all_pairs):,} 对") - - # ====== 参数 ====== - TARGET = 2000000 # 目标200万对 - BATCH = 40 # 每批40个任务 - MAX_WORKERS = 8 # 8并发（比之前5更猛） - CHECKPOINT = 5000 # 每5000对保存 - - to_generate = TARGET - len(all_pairs) - total_batches = (to_generate + BATCH - 1) // BATCH - - print(f"目标: {TARGET:,} 对") - print(f"已有: {len(all_pairs):,} 对") - print(f"还需: {to_generate:,} 对") - print(f"批次: {total_batches:,} 批 ({BATCH}任务/批, {MAX_WORKERS}并发)") - - # 策略统计 - strategy_stats = {g[0]: {'runs': 0, 'pairs': 0} for g in GENERATORS} - - t_start = time.time() - last_checkpoint = len(all_pairs) - - for batch_idx in range(total_batches): - t0 = time.time() - - # 并发 - with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor: - futures = [executor.submit(run_one_generation) for _ in range(BATCH)] - batch_pairs = [] - for future in as_completed(futures): - try: - r = future.result() - if r: - batch_pairs.extend(r) - except: - pass - - # 去重 - new_pairs = [] - for q, a in batch_pairs: - key = (q, a) - if key not in existing_keys: - existing_keys.add(key) - new_pairs.append({'question': q, 'answer': a}) - - all_pairs.extend(new_pairs) - - elapsed = time.time() - t0 - total_elapsed = time.time() - t_start - rate = len(all_pairs) / max(total_elapsed, 1) * 60 - - # 估算 - pairs_this_batch = len(new_pairs) - total_raw = len(batch_pairs) - dedup_rate = pairs_this_batch / max(total_raw, 1) * 100 - - print(f"[{batch_idx+1}/{total_batches}] " - f"+{pairs_this_batch}对(有效{total_raw}/{dedup_rate:.0f}%) " - f"累计{len(all_pairs):,}/{TARGET:,} " - f"{rate:.0f}对/分 " - f"剩余{(TARGET-len(all_pairs))/max(rate,0.1):.0f}分") - - # 定期保存 - if len(all_pairs) - last_checkpoint >= CHECKPOINT: - with open(progress_file, 'w', encoding='utf-8') as f: - json.dump(all_pairs, f, ensure_ascii=False) - last_checkpoint = len(all_pairs) - - # 也保存到数据目录一份 - merged_path = os.path.join(MODEL_DIR, 'data', 'general_ongoing.json') - with open(merged_path, 'w', encoding='utf-8') as f: - json.dump(all_pairs, f, ensure_ascii=False) - - # 速率统计 - stable_rate = len(all_pairs) / max(time.time() - t_start, 1) * 60 - print(f" 💾 保存 ({len(all_pairs):,} 对, {stable_rate:.0f}对/分)") - - # ====== 完成 ====== - total_elapsed = time.time() - t_start - - with open(progress_file, 'w', encoding='utf-8') as f: - json.dump(all_pairs, f, ensure_ascii=False) - - print(f"\n{'='*65}") - print(f"✅ 生成完成!") - print(f" 总对数: {len(all_pairs):,}") - print(f" 用时: {total_elapsed/60:.1f} 分") - print(f" 平均速率: {len(all_pairs)/max(total_elapsed,1)*60:.0f} 对/分") - print(f" 保存: {progress_file}") - print(f"{'='*65}") - -if __name__ == '__main__': - main() \ No newline at end of file diff --git a/scripts/generate_multiturn.py b/scripts/generate_multiturn.py deleted file mode 100644 index 1cdd73f..0000000 --- a/scripts/generate_multiturn.py +++ /dev/null @@ -1,346 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -多轮对话生成 — 用 Astron 造 2~5 轮小狸闲聊对话 -5并发，无限量，增量保存防丢 -""" -import json, requests, time, os, sys, random, re -from concurrent.futures import ThreadPoolExecutor, as_completed -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1) - -# ====== Astron 配置（和 api_correction 一致） ====== -ASTRON_URL = "https://maas-coding-api.cn-huabei-1.xf-yun.com/v2/chat/completions" -ASTRON_KEY = os.environ.get("ASTRON_KEY", "") -ASTRON_MODEL = "astron-code-latest" - -# ====== 小狸基础提示词 ====== -SYSTEM_PROMPT = """你是一个名叫小狸的8岁中国小女孩。 - -你的人设： -- 8岁女孩子，有点害羞但很可爱 -- 说话简单自然，符合小朋友的语气 -- 回答问题要简短，1-2句话就够了 -- 不要用"喵"结尾，不要过度卖萌 - -回答要求： -1. 语气自然，像真实的小朋友说话 -2. 不懂就说"不知道" -3. 回答简洁但要说清楚""" - -# ====== 多轮对话种子（开场话题池） ====== -SEEDS = [ - # 日常问候 - "你好呀", "你叫什么名字", "你今天开心吗", "你在干嘛呢", - "你好可爱呀", "我来看你啦", "你今天心情怎么样", - # 闲聊话题 - "你喜欢什么颜色", "你最喜欢什么动物", "你爱吃什么零食", - "你最喜欢什么玩具", "你最喜欢什么动画片", - "你最好的朋友是谁", "幼儿园里好玩吗", - "你喜欢唱歌吗", "你会跳舞吗", - # 情绪表达 - "我今天好开心", "我有点不开心", "我好无聊啊", - "我遇到了一件好事", "我今天有点难过", - "我好想你呀", "你能不能陪我聊聊天", - # 好奇问题 - "为什么天是蓝的", "为什么月亮跟着我走", - "世界上最大的动物是什么", "你会讲故事吗", - "你会数数吗", "1+1等于几", - "你喜欢上学吗", "你今年几岁啦", - # 分享日常 - "我今天去公园玩了", "我养了一只小猫", - "我今天考试考了100分", "我学会骑自行车了", - "妈妈给我买了个新玩具", "我今天摔了一跤", - # 晚安/告别 - "我要睡觉了", "明天见", "晚安啦", - "我要去吃饭了", "下次再聊哦", -] - -# 话题扩展 — 组合出更多种子 -PREFIXES = ["小狸", "小狸小狸", "亲爱的小狸", "嘿嘿小狸", "小狸呀"] -SUFFIXES = ["呀", "吗", "呢", "", "呀", "吗", "呢", ""] - -def expand_seeds(base, max_expand=5): - """种子话题 → 多样化的开场白""" - expanded = [] - for s in base: - expanded.append(s) - if '吗' not in s and not s.endswith('呀') and not s.endswith('呢'): - for _ in range(2): - prefix = random.choice(PREFIXES) if random.random() > 0.4 else '' - suffix = random.choice(SUFFIXES) if random.random() > 0.5 else '' - combo = f"{prefix}{s}{suffix}".strip() - if combo and combo not in expanded: - expanded.append(combo) - random.shuffle(expanded) - return expanded[:len(base) * max_expand] - - -def call_astron_multi(prompt, temperature=0.85): - """调用 Astron 生成一条多轮对话（2~5轮）""" - for attempt in range(3): - try: - headers = {"Content-Type": "application/json", "Authorization": f"Bearer {ASTRON_KEY}"} - payload = { - "model": ASTRON_MODEL, - "messages": [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": prompt} - ], - "temperature": temperature, - "max_tokens": 8192, - } - resp = requests.post(ASTRON_URL, headers=headers, json=payload, timeout=60) - if resp.status_code == 200: - content = resp.json()['choices'][0]['message']['content'].strip() - if content and len(content) > 10: - return content - else: - time.sleep(1) - except Exception: - time.sleep(2) - return None - - -def parse_multi_turn(text): - """ - 解析 Astron 返回的多轮对话文本为 [(q, a), (q, a), ...] - 期望返回格式（至少2轮对话）： - 问：xxx - 答：xxx - 问：xxx - 答：xxx - - 或 - 用户：xxx - 小狸：xxx - 用户：xxx - 小狸：xxx - """ - lines = text.strip().split('\n') - turns = [] - current_q = None - - for line in lines: - line = line.strip() - if not line: - continue - - # 匹配 "问：" 或 "问题：" 或 "用户：" 格式 - q_match = re.match(r'^(问|问题|用户|我)[：:]\s*(.*)', line) - # 匹配 "答：" 或 "回答：" 或 "小狸：" 格式 - a_match = re.match(r'^(答|回答|小狸|小狸说|机器人)[：:]\s*(.*)', line) - - if q_match: - q_text = q_match.group(2).strip() - if q_text and len(q_text) > 1: - current_q = q_text - elif a_match and current_q: - a_text = a_match.group(2).strip() - if a_text and len(a_text) > 1 and len(a_text) < 200: - turns.append((current_q, a_text)) - current_q = None # 配对后重置，等待下一个问 - - # 最少4轮才算有效（6~12轮的长对话） - if len(turns) >= 4: - return turns - return None - - -def format_training_sample(turns): - """ - 将多轮对话 [('q1','a1'), ('q2','a2'), ...] 格式化为训练样本 - - 格式：q1a1q2a2 - """ - parts = [] - for i, (q, a) in enumerate(turns): - if i == 0: - parts.append(f"{q}{a}") - else: - parts.append(f"{q}{a}") - parts.append("") - return "".join(parts) - - -def make_multi_prompt(): - """ - 构造多轮对话生成提示词 - Astron 每次生成一条完整的4~6轮对话 - """ - # 随机选开场话题（2~3个组合，让对话有展开） - topic_count = random.randint(1, 3) - topics = random.sample(all_seeds, min(topic_count, len(all_seeds))) - topic_context = "、".join(topics) - - turn_count = random.randint(6, 12) - - prompt = f"""请以小狸的身份，模拟一段长时间的自然闲聊对话，至少{turn_count}轮。 - -开场话题：{topic_context} - -要求： -1. 小狸是8岁小女孩，说话简短自然 -2. 每轮包含"问"和"答"，小狸的回答1-2句话 -3. 对话要自然展开，一问一答慢慢深入，不要强行切话题 -4. 小狸的语气：害羞但可爱，不是机器人 -5. 至少{turn_count}轮，越多越好，让对话持续下去 - -请按以下格式输出（至少{turn_count}轮，轮数越多越好）： -问：（用户提问） -答：（小狸回答） -问：（用户续问） -答：（小狸回答） -...""" - return prompt - - -def generate_one_multi(): - """生成一条多轮对话""" - prompt = make_multi_prompt() - result = call_astron_multi(prompt, temperature=random.uniform(0.8, 1.0)) - if result: - turns = parse_multi_turn(result) - if turns and len(turns) >= 4: - sample = format_training_sample(turns) - return { - 'source': 'multiturn_astron', - 'turns': len(turns), - 'sample': sample, - 'turns_list': turns, - 'raw': result[:100] - } - return None - - -def main(): - global all_seeds - - model_dir = os.path.dirname(os.path.abspath(__file__)) - - # 加载现有训练数据的问句作为种子 - train_path = os.path.join(model_dir, 'xiaoli_train_v2.json') - with open(train_path, encoding='utf-8') as f: - train_data = json.load(f) - - # 从现有数据中随机抽取问句作为种子 - existing_qs = [item.get('question','') for item in train_data if item.get('question','')] - random.shuffle(existing_qs) - seed_from_data = existing_qs[:500] - - # 合并种子话题 - all_seeds = expand_seeds(SEEDS, max_expand=10) + seed_from_data - random.shuffle(all_seeds) - - print(f"种子话题: {len(all_seeds)} 个 (基础{len(SEEDS)}+数据抽取{len(seed_from_data)})") - - # 输出目录 - out_dir = os.path.join(model_dir, 'multiturn_outputs') - os.makedirs(out_dir, exist_ok=True) - - # 检查已有进度 - all_results = [] - interim_files = sorted([f for f in os.listdir(out_dir) - if f.startswith('mt_batch_') and f.endswith('.json')]) - if interim_files: - for fname in interim_files: - with open(os.path.join(out_dir, fname), encoding='utf-8') as f: - data = json.load(f) - all_results.extend(data) - print(f"已有中间结果: {len(all_results)} 条 ({len(interim_files)} 批)") - - # ====== 参数 ====== - TARGET = 500000 # 目标 50 万条长多轮对话 - BATCH = 20 # 每批生成 20 条（每条更长，减少批量） - MAX_WORKERS = 5 # Astron 5并发 - TEMPERATURE_RANGE = (0.80, 1.0) - - to_generate = TARGET - len(all_results) - total_batches = (to_generate + BATCH - 1) // BATCH - - print(f"目标: {TARGET} 条, 已有: {len(all_results)} 条, 还需: {to_generate} 条") - print(f"批次: {total_batches} 批 (每批{BATCH}条, {MAX_WORKERS}并发)") - print(f"{'='*60}\n") - - t_start = time.time() - - for batch_idx in range(total_batches): - t0 = time.time() - - # 并发生成一批 - with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor: - futures = [executor.submit(generate_one_multi) for _ in range(BATCH)] - batch_results = [] - for future in as_completed(futures): - try: - r = future.result() - if r: - batch_results.append(r) - except Exception: - pass - - # 去重（按 sample 内容） - existing_samples = set(r['sample'] for r in all_results) - deduped = [] - for r in batch_results: - if r['sample'] not in existing_samples: - existing_samples.add(r['sample']) - deduped.append(r) - - all_results.extend(deduped) - - # 保存本批 - elapsed = time.time() - t0 - batch_file = os.path.join(out_dir, f'mt_batch_{batch_idx + 1:04d}.json') - with open(batch_file, 'w', encoding='utf-8') as f: - json.dump(deduped, f, ensure_ascii=False, indent=2) - - # 追加保存总进度 - progress_file = os.path.join(out_dir, 'all_multiturn_progress.json') - with open(progress_file, 'w', encoding='utf-8') as f: - json.dump(all_results, f, ensure_ascii=False, indent=2) - - total_elapsed = time.time() - t_start - rate = len(all_results) / max(total_elapsed, 1) * 60 - eta = (TARGET - len(all_results)) / max(rate, 0.1) / 60 if rate > 0 else 0 - - # 输出统计 - avg_turns = sum(r['turns'] for r in deduped) / max(len(deduped), 1) - print(f"批次 {batch_idx+1}/{total_batches} | " - f"+{len(deduped)} 条 (去重后) | " - f"累计 {len(all_results)}/{TARGET} | " - f"均轮 {avg_turns:.1f} | " - f"{rate:.0f}条/分 | " - f"预计剩余 {eta:.1f}分") - - # ====== 完成总结 ====== - total_elapsed = time.time() - t_start - print(f"\n{'='*60}") - print(f"多轮对话生成完成!") - print(f"生成: {len(all_results)} 条 (去重后)") - avg_turns2 = sum(r['turns'] for r in all_results) / max(len(all_results), 1) - print(f"平均轮数: {avg_turns2:.1f}") - print(f"总对话轮次: {sum(r['turns'] for r in all_results)}") - print(f"用时: {total_elapsed/60:.1f} 分钟") - print(f"速率: {len(all_results)/max(total_elapsed,1)*60:.0f} 条/分") - - # ====== 保存最终结果（独立存放，不追加到训练集） ====== - final_file = os.path.join('/data/user/work', 'multiturn_final.json') - with open(final_file, 'w', encoding='utf-8') as f: - json.dump(all_results, f, ensure_ascii=False, indent=2) - print(f"\n最终结果已保存: {final_file} ({len(all_results)} 条)") - print(f"如需后续使用，可手动合并到训练集") - - # ====== 保存多轮样本供训练时直接使用（含标记格式） ====== - mt_samples_file = os.path.join(model_dir, 'multiturn_samples_v1.json') - with open(mt_samples_file, 'w', encoding='utf-8') as f: - samples = [{'turns': r['turns'], 'sample': r['sample']} for r in all_results] - json.dump(samples, f, ensure_ascii=False, indent=2) - print(f" ✅ 多轮样本保存: multiturn_samples_v1.json ({len(samples)} 条)") - - print(f"\n{'='*60}") - print(f"全部完成! 新增 {len(all_results)} 条多轮对话") - print(f"{'='*60}") - -if __name__ == '__main__': - main() \ No newline at end of file diff --git a/scripts/generate_real_multiturn.py b/scripts/generate_real_multiturn.py deleted file mode 100644 index 4811229..0000000 --- a/scripts/generate_real_multiturn.py +++ /dev/null @@ -1,337 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -海量真实多轮对话生成 — 模拟真实用户，8~15轮长对话 -""" -import json, requests, time, os, sys, random, re -from concurrent.futures import ThreadPoolExecutor, as_completed -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1) - -from scripts.api_correction import batch_generate - -MODEL_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - -# ====== Astron 配置（复用 api_correction 的 key） ====== -ASTRON_URL = "https://maas-coding-api.cn-huabei-1.xf-yun.com/v2/chat/completions" -ASTRON_KEY = os.environ.get("ASTRON_KEY", "") -ASTRON_MODEL = "astron-code-latest" - -# ====== 用户场景画像（模拟真实用户） ====== -USER_SCENARIOS = [ - # 小小孩子（3-5岁） - {"age": "3-5岁", "style": "说话简单，喜欢问为什么，爱说叠词", "tone": "奶声奶气"}, - # 小孩子（6-9岁） - {"age": "6-9岁", "style": "小学生，爱分享学校的事，喜欢动画片和游戏", "tone": "活泼天真"}, - # 大孩子（10-14岁） - {"age": "10-14岁", "style": "有点小叛逆，有自己的想法，会吐槽", "tone": "有点酷"}, - # 妈妈/家长 - {"age": "成年", "style": "关心孩子，会问小狸吃饭睡觉学习的事", "tone": "温柔耐心"}, - # 孤单的人 - {"age": "成年", "style": "有点孤独，需要陪伴，会把小狸当倾诉对象", "tone": "低落温柔"}, - # 高中生/大学生 - {"age": "16-22岁", "style": "有压力，会聊学业、未来、人际关系", "tone": "成熟但需要安慰"}, - # 好奇宝宝 - {"age": "不限", "style": "什么都想知道，十万个为什么", "tone": "充满好奇"}, - # 伤心的人 - {"age": "不限", "style": "心情不好，需要安慰和陪伴", "tone": "低沉脆弱"}, -] - -# ====== 真实对话种子（更自然的开场） ====== -REAL_SEEDS = [ - # 日常打招呼（多种自然说法） - "小狸你在吗", "嘿嘿小狸", "小狸小狸快出来", "小狸你在干嘛呀", - "我来找你玩啦", "小狸今天乖不乖", "你在做什么呢", - # 情绪分享（真实用户会说的话） - "我今天好不开心", "我开心死了哈哈哈", "我好无聊啊没人陪我", - "今天倒霉死了", "我今天遇到了好事", "我有点难过", - "我心情不好你能陪我吗", "我好想你啊", "我好孤单", - "我害怕", "我有点紧张明天要考试", "我好烦", - # 日常分享 - "今天我妈做了红烧肉", "我们家养了一只猫超可爱", "我今天考试考了98分", - "我感冒了好难受", "我今天被老师表扬了", "我和朋友吵架了", - "我昨天做梦梦到你了", "我学会骑自行车了！", "我今天摔了一跤好疼", - "学校要考试了好紧张", "放假了好开心", "又要开学了好痛苦", - "今天下雨了不能出去玩", "今天天气好好我们去玩吧", - "我妈妈给我买了个新玩具", "我养的花开了", - # 求助/询问 - "我不小心把杯子打碎了怎么办", "我作业好多写不完", - "我睡不着你能陪我聊天吗", "我牙好疼", - "妈妈骂我了", "我被老师批评了", - "我迷路了怎么办", "我找不到我的水杯了", - # 好奇/知识 - "为什么天是蓝的呀", "鱼会睡觉吗", "为什么会有影子", - "星星为什么会闪", "海为什么是咸的", "为什么要上学", - "为什么大人总说不行", "什么是爱", "什么是长大", - # 表达好感的 - "小狸我最喜欢你了", "你好可爱呀", "我们是最好的朋友对吧", - "我送给你一个礼物", "小狸你真聪明", "你真好", - # 关于小狸本身的 - "你几岁啦", "你喜欢什么", "你有爸爸妈妈吗", - "你住在哪里呀", "你喜欢吃什么", "你最好的朋友是谁", - "你上学吗", "你的梦想是什么", "你平时喜欢做什么", - # 奇怪/无厘头 - "如果我能飞就好了", "我想变成一只猫", "你觉得外星人存在吗", - "假如你会魔法你想做什么", "你相信世界上有奥特曼吗", - # 夜间/睡前 - "我睡不着", "陪我聊到睡着好不好", "给我讲个睡前故事", - "晚安前陪我说说话", "我怕黑", - # 亲情家庭 - "我爸妈吵架了怎么办", "我想我奶奶了", "我好久没见爸爸了", - "妈妈生病了", "我弟弟好烦总是抢我东西", - # 节假日 - "过年啦新年快乐", "今天是我生日", "圣诞节到了", - "六一儿童节快乐", "明天要去春游好开心", -] - -CONTEXTS = [ - "晚上睡前躺在床上", "放学回家的路上", - "周末在家没事做", "刚吃完饭坐着聊天", - "在公园里坐着", "下雨天窝在家里", - "一个人在房间", "在院子里玩的时候", - "刚写完作业休息", "早上起床的时候", - "在操场上看别人玩", "坐车去玩的路上", - "在图书馆里", "在外面吃饭等上菜", - "一个人发呆的时候", "看着窗外的时候", -] - -def make_realistic_multi_prompt(): - """构造非常逼真的长多轮对话提示 — 一次生成4条""" - parts = [] - for _ in range(4): - scenario = random.choice(USER_SCENARIOS) - opener = random.choice(REAL_SEEDS) - context = random.choice(CONTEXTS) - min_turns = random.randint(6, 8) - max_turns = random.randint(10, 14) - - parts.append(f"""=== 对话 {_+1} === -场景：{context} -用户：{scenario['age']}，{scenario['style']}，{scenario['tone']} -开场：{opener} -对话： -用户：{opener} -小狸：""") - - prompt = f"""请生成 {len(parts)} 段独立的日常聊天对话，每段 {min_turns}~{max_turns} 轮。 - -重要：每段对话必须是一个完整的聊天，话题自然展开，不要强行结束。 - -{"".join(parts)} - -严格要求： -1. 每个=== 对话 ===之后的段落实例必须写完整，不要留空 -2. 用户提问必须口语化、自然，像真人在说话 -3. 小狸回答1-2句话，像8岁小朋友 -4. 话题要延续上一轮，逐渐深入延伸 -5. 用户要有真实反应（惊讶、追问、分享等） -6. 用"问："/"答："格式 -7. 输出不能省略任何一轮 -""" - return prompt - -def call_astron_multi(prompt, temperature=0.85): - for attempt in range(3): - try: - headers = {"Content-Type": "application/json", "Authorization": f"Bearer {ASTRON_KEY}"} - payload = { - "model": ASTRON_MODEL, - "messages": [ - {"role": "system", "content": "你是一个擅长写真实对话的助手。"}, - {"role": "user", "content": prompt} - ], - "temperature": temperature, - "max_tokens": 8192, - } - resp = requests.post(ASTRON_URL, headers=headers, json=payload, timeout=120) - if resp.status_code == 200: - content = resp.json()['choices'][0]['message']['content'].strip() - if content and len(content) > 50: - return content - else: - time.sleep(1) - except Exception: - time.sleep(2) - return None - -def parse_multi_turn(text): - """解析多段对话（支持=== 对话 N ===分隔，也支持纯轮次）""" - lines = text.strip().split('\n') - conversations = [] - current_q = None - current_turns = [] - - for line in lines: - line = line.strip() - if not line: - continue - - # 检测新对话开始标记 - if re.match(r'===?\s*对话\s*\d+\s*===?', line): - if current_turns and len(current_turns) >= 4: - conversations.append(current_turns) - current_turns = [] - current_q = None - continue - - # 跳过设定行（场景/用户/开场） - if line.startswith(('场景', '用户', '开场', '对话', '===')): - if line.startswith(('场景', '用户', '开场')): - continue - continue - - # 用户提问 - q_match = re.match(r'^(用户|我|问|问题|小孩|小朋友|孩子|朋友)\s*[：:]\s*(.*)', line) - a_match = re.match(r'^(小狸|答|回答|机器人)\s*[：:]\s*(.*)', line) - - if q_match: - q_text = q_match.group(2).strip() - if q_text and len(q_text) > 1: - current_q = q_text - elif a_match and current_q: - a_text = a_match.group(2).strip() - if a_text and len(a_text) > 1 and len(a_text) < 300: - current_turns.append((current_q, a_text)) - current_q = None - - # 保存最后一组 - if current_turns and len(current_turns) >= 4: - conversations.append(current_turns) - - # 过滤：只保留>=4轮的对话，且累积轮次>=8 - valid = [c for c in conversations if len(c) >= 4] - total_turns = sum(len(c) for c in valid) - if total_turns >= 8 and valid: - return valid - return None - -def make_turns_list_text(turns): - """将turns转成用户可见的对话文本""" - lines = [] - for i, (q, a) in enumerate(turns): - lines.append(f"问：{q}") - lines.append(f"答：{a}") - return "\n".join(lines) - -def generate_one_multi(): - """生成多条多轮对话（一次API调用返回多条）""" - prompt = make_realistic_multi_prompt() - result = call_astron_multi(prompt, temperature=random.uniform(0.75, 0.95)) - if result: - convs = parse_multi_turn(result) - if convs: - items = [] - for turns in convs: - items.append({ - 'turns': len(turns), - 'turns_list': turns, - 'text': make_turns_list_text(turns), - 'source': 'astron_multi_v2', - }) - return items - return None - -def main(): - print("="*60) - print(" 海量真实多轮对话生成 v2") - print(" 生成超长、超真实的多轮聊天") - print("="*60) - - out_dir = os.path.join(MODEL_DIR, 'multiturn_outputs') - os.makedirs(out_dir, exist_ok=True) - - # 检查已有进度 - all_results = [] - progress_file = os.path.join(out_dir, 'all_multiturn_progress.json') - if os.path.exists(progress_file): - all_results = json.load(open(progress_file, encoding='utf-8')) - print(f"已有进度: {len(all_results)} 条") - - # ====== 参数 ====== - TARGET = 500000 # 目标50万条 - BATCH = 30 # 每批30条 - MAX_WORKERS = 5 # Astron 5并发（稳定不超限） - CHECKPOINT = 1000 # 每1000条保存一次进度 - - to_generate = TARGET - len(all_results) - total_batches = (to_generate + BATCH - 1) // BATCH - - print(f"目标: {TARGET:,} 条") - print(f"已有: {len(all_results):,} 条") - print(f"还需: {to_generate:,} 条") - print(f"批次: {total_batches:,} 批 (每批{BATCH}条, {MAX_WORKERS}并发)") - print(f"{'='*60}\n") - - t_start = time.time() - last_checkpoint = len(all_results) - - for batch_idx in range(total_batches): - t0 = time.time() - - # 并发生成 - with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor: - futures = [executor.submit(generate_one_multi) for _ in range(BATCH)] - batch_results = [] - for future in as_completed(futures): - try: - r = future.result() - if r: - batch_results.extend(r) - except Exception as e: - pass - - # 去重 - existing_texts = set(r.get('text', '')[:200] for r in all_results) - deduped = [] - for r in batch_results: - key = r.get('text','')[:200] - if key not in existing_texts: - existing_texts.add(key) - deduped.append(r) - - all_results.extend(deduped) - - elapsed = time.time() - t0 - total_elapsed = time.time() - t_start - rate = len(all_results) / max(total_elapsed, 1) * 60 - eta = (TARGET - len(all_results)) / max(rate, 0.1) / 60 if rate > 0 else 0 - - # 统计 - valid = len(batch_results) - avg_turns = sum(r['turns'] for r in deduped) / max(len(deduped), 1) - success_rate = valid / BATCH * 100 - - print(f"[{batch_idx+1}/{total_batches}] " - f"+{len(deduped)}条(有效{valid}/{success_rate:.0f}%) " - f"累计{len(all_results):,}/{TARGET:,} " - f"均{avg_turns:.1f}轮 " - f"{rate:.0f}条/分 " - f"剩余{eta:.0f}分") - - # 定期保存进度 - if len(all_results) - last_checkpoint >= CHECKPOINT: - with open(progress_file, 'w', encoding='utf-8') as f: - json.dump(all_results, f, ensure_ascii=False) - last_checkpoint = len(all_results) - print(f" 💾 已保存进度 ({len(all_results):,} 条)") - - # ====== 完成 ====== - total_elapsed = time.time() - t_start - with open(progress_file, 'w', encoding='utf-8') as f: - json.dump(all_results, f, ensure_ascii=False) - - avg_turns_all = sum(r['turns'] for r in all_results) / max(len(all_results), 1) - total_turns = sum(r['turns'] for r in all_results) - - print(f"\n{'='*60}") - print(f"✅ 生成完成!") - print(f" 总条数: {len(all_results):,}") - print(f" 平均轮数: {avg_turns_all:.1f}") - print(f" 总对话轮次: {total_turns:,}") - print(f" 用时: {total_elapsed/60:.1f} 分") - print(f" 速率: {len(all_results)/max(total_elapsed,1)*60:.0f} 条/分") - print(f" 保存: {progress_file}") - print(f"{'='*60}") - -if __name__ == '__main__': - main() \ No newline at end of file diff --git a/scripts/regenerate_all_high_temp.py b/scripts/regenerate_all_high_temp.py deleted file mode 100644 index d97214f..0000000 --- a/scripts/regenerate_all_high_temp.py +++ /dev/null @@ -1,126 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -高温多样化生成 — 把训练集全部问题用 Astron 高温度重答一次 -温度=1.2，追求多样化回答，丰富训练数据 -""" -import sys, os, json, time - -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1) - -from api_correction import batch_generate - -MODEL_DIR = os.path.dirname(os.path.abspath(__file__)) -TEMPERATURE = 1.0 # 高温度(上限1.0)，产出多样化回答 -BATCH_SIZE = 500 # 每批 500 条，分批保存防丢 -MAX_WORKERS = 5 # Astron 5并发 - -def extract_questions(data): - """从训练数据中提取所有问题""" - questions = [] - for item in data: - if isinstance(item, dict): - q = item.get('question', '') or item.get('input', '') - if q: - questions.append(q) - return questions - -def main(): - # 1. 加载训练数据 - train_path = os.path.join(MODEL_DIR, 'xiaoli_train_v2.json') - print(f"加载训练数据: {train_path}") - with open(train_path, encoding='utf-8') as f: - train_data = json.load(f) - - questions = extract_questions(train_data) - total = len(questions) - print(f"提取到 {total} 条问题，温度={TEMPERATURE}，{MAX_WORKERS}并发") - - # 2. 创建输出目录 - out_dir = os.path.join(MODEL_DIR, 'high_temp_outputs') - os.makedirs(out_dir, exist_ok=True) - - # 3. 分批处理 - all_results = [] - existing = set() # 去重用 - - # 检查是否有已完成的中间结果 - interim_files = sorted([f for f in os.listdir(out_dir) if f.startswith('batch_') and f.endswith('.json')]) - if interim_files: - for fname in interim_files: - with open(os.path.join(out_dir, fname), encoding='utf-8') as f: - batch_data = json.load(f) - for r in batch_data: - key = (r['question'][:100], r['answer'][:100]) - if key not in existing: - existing.add(key) - all_results.append(r) - print(f"发现已有中间结果: {len(all_results)} 条（来自 {len(interim_files)} 个文件）") - # 跳过已处理的 question - processed_questions = set(r['question'] for r in all_results) - remaining = [q for q in questions if q not in processed_questions] - print(f"已处理 {len(processed_questions)} 条，剩余 {len(remaining)} 条") - questions = remaining - - total_batches = (len(questions) + BATCH_SIZE - 1) // BATCH_SIZE - t_start = time.time() - - for batch_idx in range(total_batches): - batch_qs = questions[batch_idx * BATCH_SIZE : (batch_idx + 1) * BATCH_SIZE] - print(f"\n{'='*60}") - print(f"批次 {batch_idx + 1}/{total_batches} ({len(batch_qs)}条, 温度{TEMPERATURE})") - print(f"{'='*60}") - - t0 = time.time() - results = batch_generate(batch_qs, temperature=TEMPERATURE, max_workers=MAX_WORKERS) - elapsed = time.time() - t0 - - # 标注入参温度 - for r in results: - r['temperature'] = TEMPERATURE - - # 去重后加入总集 - new_in_batch = 0 - for r in results: - key = (r['question'][:100], r['answer'][:100]) - if key not in existing: - existing.add(key) - all_results.append(r) - new_in_batch += 1 - - # 立即保存本批结果 - batch_file = os.path.join(out_dir, f'batch_{batch_idx + 1:04d}.json') - with open(batch_file, 'w', encoding='utf-8') as f: - json.dump(results, f, ensure_ascii=False, indent=2) - - # 追加保存总进度 - progress_file = os.path.join(out_dir, 'all_progress.json') - with open(progress_file, 'w', encoding='utf-8') as f: - json.dump(all_results, f, ensure_ascii=False, indent=2) - - total_elapsed = time.time() - t_start - rate = len(all_results) / max(total_elapsed, 1) * 60 - eta_remaining = (total - len(all_results)) / max(rate, 0.1) / 60 if rate > 0 else 0 - print(f" 本批新增 {new_in_batch} 条 (去重后) | 累计 {len(all_results)} 条 | " - f"{rate:.0f}条/分 | 已用 {total_elapsed/60:.1f}分 | 预计剩余 {eta_remaining:.1f}分") - - # 4. 去重统计 - print(f"\n{'='*60}") - print(f"全部生成完成!") - print(f"生成结果: {len(all_results)} 条 (去重后)") - print(f"用时: {(time.time() - t_start)/60:.1f} 分钟") - - # 5. 保存最终去重结果（独立存放，不追加到训练集） - final_file = os.path.join('/data/user/work', 'high_temp_final.json') - with open(final_file, 'w', encoding='utf-8') as f: - json.dump(all_results, f, ensure_ascii=False, indent=2) - print(f"\n最终结果已保存: {final_file}") - print(f"如需使用，可手动合并到训练集") - - print(f"\n{'='*60}") - print(f"全部完成! 新增 {len(all_results)} 条多样化回答") - print(f"{'='*60}") - -if __name__ == '__main__': - main() \ No newline at end of file diff --git a/scripts/train_cnb.py b/scripts/train_cnb.py deleted file mode 100644 index 75e7c17..0000000 --- a/scripts/train_cnb.py +++ /dev/null @@ -1,204 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -小狸AI - CNB GPU训练脚本 -从 merged_clean_v3.json (977K对) 训练，自动GPU检测 -""" -import sys, os, json, time, glob, re, random, argparse -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1) - -import torch -import torch.nn as nn -from torch.utils.data import Dataset, DataLoader - -from llm_chat_improved import LLMModel, LLMChat, ConversationDataset, LLMTrainer, tokenizer_instance - -MODEL_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - -def load_pairs(path): - with open(path, 'r', encoding='utf-8') as f: - data = json.load(f) - pairs = [] - for item in data: - if isinstance(item, dict): - q = item.get("question", "") or item.get("input", "") - a = item.get("answer", "") or item.get("output", "") - if q and a: - pairs.append((q, a)) - return pairs - -def find_latest_ckpt(model_dir): - ckpts = glob.glob(os.path.join(model_dir, 'llm_model_v2_resume_epoch_*.pth')) - if not ckpts: - ckpts = glob.glob(os.path.join(model_dir, 'llm_model_v2_epoch_*.pth')) - if not ckpts: - return None, 0 - nums = [(p, max([int(s) for s in re.findall(r'\d+', os.path.basename(p).replace('.pth',''))])) for p in ckpts] - best = max(nums, key=lambda x: x[1]) - return best[0], best[1] - -def main(): - parser = argparse.ArgumentParser(description='小狸AI CNB训练') - parser.add_argument('--epochs', type=int, default=100, help='训练轮数') - parser.add_argument('--batch_size', type=int, default=8, help='batch大小') - parser.add_argument('--lr', type=float, default=3e-4, help='学习率') - parser.add_argument('--accum', type=int, default=2, help='梯度累积步数') - parser.add_argument('--data', type=str, default='data/merged_clean_v3.json', help='训练数据路径') - parser.add_argument('--resume', action='store_true', help='断点续训') - parser.add_argument('--save_every', type=int, default=5, help='每N轮保存checkpoint') - args = parser.parse_args() - - device = 'cuda' if torch.cuda.is_available() else 'cpu' - print("=" * 60) - print(f"小狸AI CNB训练") - print(f" GPU: {torch.cuda.get_device_name(0) if device == 'cuda' else '无'}") - print(f" 设备: {device}") - print(f" 数据: {args.data}") - print(f" epochs={args.epochs}, batch={args.batch_size}, lr={args.lr}") - print("=" * 60) - - # ====== 加载数据 ====== - data_path = args.data if os.path.isabs(args.data) else os.path.join(MODEL_DIR, args.data) - if not os.path.exists(data_path): - print(f"❌ 数据不存在: {data_path}") - return - - pairs = load_pairs(data_path) - random.seed(42) - random.shuffle(pairs) - split = int(len(pairs) * 0.95) - train_set = pairs[:split] - val_set = pairs[split:] - print(f"数据: {len(pairs)}对 | 训练: {len(train_set)} | 验证: {len(val_set)}") - - # ====== 加载词表 ====== - vocab_path = os.path.join(MODEL_DIR, 'vocab.json') - with open(vocab_path, 'r', encoding='utf-8') as f: - vocab = json.load(f) - # vocab.json 不含特殊token(0-6)，需合并进来，否则： - # 1) tokenizer 编码 // 等标记时会逐字符错码 - # 2) vocab 的 id 范围是 7..max，超过 len(vocab) 会造成 embedding 索引越界 - special_tokens = {'': 0, '': 1, '': 2, '': 3, - '': 4, '': 5, '': 6} - for tok, tid in special_tokens.items(): - if tok not in vocab: - vocab[tok] = tid - real_vocab_size = max(vocab.values()) + 1 - print(f"词表: {len(vocab)}项, vocab_size={real_vocab_size} (字符+BPE, 来自 vocab.json)") - tokenizer_instance.char_to_id = vocab.copy() - tokenizer_instance.id_to_char = {v: k for k, v in vocab.items()} - tokenizer_instance.special_tokens = special_tokens.copy() - tokenizer_instance._sorted_vocab_cache = None - for attr in ['_multi_trie', '_multi_keys', '_single_map']: - if hasattr(tokenizer_instance, attr): - delattr(tokenizer_instance, attr) - - # ====== 模型 ====== - model = LLMModel(vocab_size=real_vocab_size, d_model=320, n_heads=8, - n_layers=6, d_ff=1536, max_len=256, dropout=0.1, - use_memory=False, use_thinking=False, use_moe=False) - total = sum(p.numel() for p in model.parameters()) - print(f"模型: {total:,} 参数 ({total/1e6:.2f}M)") - - start_epoch = 0 - if args.resume: - ckpt_path, last_epoch = find_latest_ckpt(MODEL_DIR) - if ckpt_path: - ckpt = torch.load(ckpt_path, map_location=device) - old_cfg = ckpt.get('model_config', {}) - if (old_cfg.get('vocab_size', 0) == real_vocab_size - and old_cfg.get('d_model', 0) == 320 - and old_cfg.get('n_layers', 0) == 6): - model.load_state_dict(ckpt['model_state_dict']) - start_epoch = last_epoch - print(f"断点续训: epoch {start_epoch} ← {os.path.basename(ckpt_path)}") - else: - print(f"⚠️ 架构不匹配，从头训练") - else: - print("无断点，从头训练") - - model = model.to(device) - - # ====== 火力全开：torch.compile (inductor融合, 无CUDA graphs) + 梯度累积 ====== - # 物理batch=192(不OOM), accum=2 → 等效batch=384 - fwd_model = model - if device == 'cuda': - print("🔥 启用 torch.compile (mode=default, inductor融合)...") - fwd_model = torch.compile(model, mode='default') - print("✓ torch.compile 已就绪（首步触发编译约30-60s）") - - # ====== DataLoader ====== - n_workers = 16 # CPU火力全开，16核全用 - train_ds = ConversationDataset(train_set, vocab, max_length=256) - train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, - num_workers=n_workers, pin_memory=(device == 'cuda'), - persistent_workers=True, prefetch_factor=4, drop_last=True) - val_ds = ConversationDataset(val_set, vocab, max_length=256) - val_loader = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False, - num_workers=n_workers, pin_memory=(device == 'cuda'), - persistent_workers=True, prefetch_factor=4, drop_last=True) - print(f"批次: 训练{len(train_loader)} | 验证{len(val_loader)} | num_workers={n_workers} | prefetch=4 | drop_last=True") - - # ====== Trainer ====== - config = { - 'batch_size': args.batch_size, - 'gradient_accumulation_steps': args.accum, - 'epochs': args.epochs, - 'd_model': 320, 'n_heads': 8, 'n_layers': 6, - 'd_ff': 1536, 'max_length': 256, 'dropout': 0.1, - 'use_amp': device == 'cuda', - } - trainer = LLMTrainer(fwd_model, device, lr=args.lr, use_amp=config['use_amp']) - - # ====== 训练 ====== - print(f"\n开始训练 epoch {start_epoch+1} → {args.epochs}") - print("=" * 60) - t_start = time.time() - best_val = float('inf') - - for epoch in range(start_epoch + 1, args.epochs + 1): - t0 = time.time() - train_loss = trainer.train_epoch(train_loader, epoch, vocab, config) - val_loss = trainer.evaluate(val_loader, config) - dt = time.time() - t0 - - marker = " ★" if val_loss < best_val else "" - if val_loss < best_val: - best_val = val_loss - trainer.save_model(os.path.join(MODEL_DIR, 'xiaoli_cnb_best.pth'), vocab, epoch) - - elapsed = time.time() - t_start - eta = (elapsed / (epoch - start_epoch)) * (args.epochs - epoch) if epoch > start_epoch else 0 - - print(f"E{epoch}/{args.epochs} | train={train_loss:.4f} val={val_loss:.4f}{marker} | " - f"{dt:.0f}s/epoch | {elapsed/60:.1f}min已用 ETA{eta/60:.0f}min") - - # 保存 - if epoch % args.save_every == 0: - ckpt_path = os.path.join(MODEL_DIR, f'llm_model_v2_resume_epoch_{epoch}.pth') - trainer.save_model(ckpt_path, vocab, epoch) - # 清理旧resume - for f in glob.glob(os.path.join(MODEL_DIR, 'llm_model_v2_resume_*.pth')): - en = max([int(s) for s in re.findall(r'\d+', os.path.basename(f).replace('.pth',''))]) - if en <= epoch - 15: - try: os.remove(f) - except: pass - print(f" 💾 已保存 epoch {epoch}") - - # ====== 完成 ====== - final_path = os.path.join(MODEL_DIR, 'xiaoli_cnb_final.pth') - trainer.save_model(final_path, vocab, args.epochs) - total_time = (time.time() - t_start) / 60 - print(f"\n✅ 训练完成! 耗时 {total_time:.1f}min") - - # 测试 - print("\n测试:") - chat = LLMChat(model, vocab, device, max_length=256) - for q in ["你好", "你是谁", "你叫什么名字", "今天天气真好", "我有点难过"]: - r = chat.generate_response(q, max_length=60, temperature=0.3) - print(f" Q: {q}") - print(f" A: {r}\n") - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/setup_cloud.sh b/setup_cloud.sh deleted file mode 100644 index 99882c0..0000000 --- a/setup_cloud.sh +++ /dev/null @@ -1,90 +0,0 @@ -#!/bin/bash -# -*- coding: utf-8 -*- -""" -XiaLi 云 GPU 一键搭建脚本 -========================= -适用于 AutoDL / 矩池云 / 恒源云 等国内云GPU平台 - -用法: - bash setup_cloud.sh # 完整搭建 - bash setup_cloud.sh --train-only # 只拉代码，不下载数据 -""" - -set -e - -echo "========================================" -echo " XiaLi 云 GPU 训练环境搭建" -echo "========================================" - -REPO_URL="https://gitee.com/shuiyu1123/xiaoli-llm-chat.git" -WORK_DIR="/root/xiaoli-llm-chat" - -# 1. 克隆代码 -echo "" -echo "[1/4] 克隆代码..." -if [ -d "$WORK_DIR" ]; then - echo " 目录已存在，拉取更新..." - cd "$WORK_DIR" && git pull -else - git clone "$REPO_URL" "$WORK_DIR" - cd "$WORK_DIR" -fi - -# 2. 安装依赖 -echo "" -echo "[2/4] 安装依赖..." -pip install torch numpy -q -# 检查 torch 是否有 CUDA -python3 -c "import torch; print(f'CUDA可用: {torch.cuda.is_available()}, GPU数: {torch.cuda.device_count()}')" 2>/dev/null || { - echo " ⚠️ 检测到无 CUDA，尝试安装 GPU 版 PyTorch..." - pip install torch --index-url https://download.pytorch.org/whl/cu118 -q -} -pip install modelscope -q - -# 确认词表存在 -echo "" -echo "[3/4] 检查词表..." -ls -la vocab.json 2>/dev/null || { - echo " ⚠️ vocab.json 不存在！请确认仓库完整" - exit 1 -} -python3 -c " -import json -with open('vocab.json') as f: - v = json.load(f) -print(f' 词表: {len(v)} tokens, vocab_size={max(v.values())+1}') -" - -# 4. 确认训练脚本 -echo "" -echo "[4/4] 确认训练脚本..." -ls -la train_cloud.py && echo " ✅ train_cloud.py 就绪" - -# 显存信息 -echo "" -echo "========================================" -echo " 环境就绪！" -echo "========================================" -python3 -c " -import torch -if torch.cuda.is_available(): - total = torch.cuda.get_device_properties(0).total_mem / 1024**3 - name = torch.cuda.get_device_name(0) - print(f' GPU: {name} ({total:.1f} GiB)') -else: - print(' ⚠️ 未检测到 GPU！') -" -echo "" -echo "启动训练命令:" -echo " cd $WORK_DIR" -echo "" -echo " # 单卡训练（完整流程）" -echo " python3 train_cloud.py --epochs 5 --align-epochs 20" -echo "" -echo " # 从检查点续训" -echo " python3 train_cloud.py --resume --epochs 10" -echo "" -echo " # 24G显存用更小batch" -echo " python3 train_cloud.py --batch-size 32" -echo "" -echo "========================================" \ No newline at end of file diff --git a/templates/chat.html b/templates/chat.html deleted file mode 100644 index 4580776..0000000 --- a/templates/chat.html +++ /dev/null @@ -1,309 +0,0 @@ - - - - - - 小狸 - AI 聊天 - - - -
-
-

🐱 小狸

-

你的 AI 小伙伴

-
-
-
-
🐱
-
你好！我是小狸，很高兴见到你！有什么想聊的吗？
-
-
-
- - -
-
-
-
🐱
-

小狸已就绪！

-

你的 AI 小伙伴正在等你聊天~

- -
-
-
- - - - diff --git a/train.py b/train.py deleted file mode 100644 index 95a57a8..0000000 --- a/train.py +++ /dev/null @@ -1,495 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -小狸完整数据集训练 - 支持断点续训 + batch=4防OOM -""" -import sys, os -sys.path.insert(0, '/workspace/xiaoli-llm-chat') -sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1) -sys.stderr = os.fdopen(sys.stderr.fileno(), 'w', buffering=1) - -import torch -import torch.nn as nn -import json -import time -import glob -import random -from torch.utils.data import Dataset, DataLoader - -from llm_chat_improved import ( - CharTokenizer, build_vocabulary, LLMModel, LLMChat, - ConversationDataset, LLMTrainer, tokenizer_instance -) -from api_correction import batch_convert, batch_generate_multi_temp - -def dict_to_tuple(data_list): - result = [] - for item in data_list: - if isinstance(item, dict): - q = item.get("question", "") or item.get("input", "") or item.get("prompt", "") or item.get("instruction", "") - a = item.get("answer", "") or item.get("output", "") or item.get("response", "") - if q and a: - result.append((q, a)) - elif isinstance(item, (list, tuple)) and len(item) >= 2: - result.append((str(item[0]), str(item[1]))) - return result - -def load_multiturn_data(path): - with open(path, 'r', encoding='utf-8') as f: - data = json.load(f) - single_turn = [] - for item in data: - if isinstance(item, dict) and "conversations" in item: - convs = item["conversations"] - for i in range(0, len(convs) - 1, 2): - if i + 1 < len(convs): - c1, c2 = convs[i], convs[i+1] - if isinstance(c1, dict): - h = c1.get("value", "") or c1.get("content", "") - a = c2.get("value", "") or c2.get("content", "") - elif isinstance(c1, list) and len(c1) >= 2: - h, a = str(c1[1]), str(c2[1]) - else: - continue - if h and a: - single_turn.append((h, a)) - return single_turn - -def find_latest_checkpoint(model_dir): - """查找最新的checkpoint，返回 (路径, epoch数) 或 (None, 0)""" - checkpoints = glob.glob(os.path.join(model_dir, 'llm_model_v2_resume_epoch_*.pth')) - if not checkpoints: - checkpoints = glob.glob(os.path.join(model_dir, 'llm_model_v2_epoch_*.pth')) - if not checkpoints: - return None, 0 - best_path = None - best_epoch = 0 - for path in checkpoints: - import re - nums = [int(s) for s in re.findall(r'\d+', os.path.basename(path).replace('.pth',''))] - if nums: - epoch = max(nums) - if epoch > best_epoch: - best_epoch = epoch - best_path = path - return best_path, best_epoch - -def migrate_weights(old_sd, new_sd_init, old_d_model, new_d_model, - old_n_layers, new_n_layers, old_d_ff, new_d_ff, - old_vocab_size, new_vocab_size): - """权重扩展迁移：将旧模型权重扩展到新架构 - 支持 d_model / n_layers / d_ff / vocab_size 全部变化 - """ - import torch - new_sd = {} - old_keys = list(old_sd.keys()) - new_keys = list(new_sd_init.keys()) - - for key in new_keys: - if 'freq' in key: - # RoPE频率表，新模型自动生成，跳过 - new_sd[key] = new_sd_init[key].clone() - continue - - # 解析层索引（如 transformer_blocks.3.attention.wq.weight） - old_key = key - new_w = new_sd_init[key].clone() - old_w_shape = None - - if key in old_sd: - old_w = old_sd[key] - old_w_shape = old_w.shape - new_shape = new_w.shape - - # 1. 严格匹配 = 直接复制 - if old_w.shape == new_shape: - new_w = old_w.clone() - # 2. 二维权重（线性层 + embedding） - elif old_w.dim() == 2: - o0, o1 = old_w.shape - n0, n1 = new_shape - # embedding.weight 和 fc_out.weight: 词汇表维度变化 - if 'embedding' in key and old_vocab_size != new_vocab_size: - # 词汇表扩展：复制旧词嵌入，新词随机 - min_v = min(old_vocab_size, new_vocab_size) - min_d = min(o1, n1) - new_w[:min_v, :min_d] = old_w[:min_v, :min_d] - print(f" 迁移 {key}: ({o0},{o1})→({n0},{n1}) vocab扩展") - elif 'fc_out' in key and old_vocab_size != new_vocab_size: - min_vocab = min(old_vocab_size, new_vocab_size) - min_feat = min(o1, n1) - new_w[:min_vocab, :min_feat] = old_w[:min_vocab, :min_feat] - print(f" 迁移 {key}: ({o0},{o1})→({n0},{n1}) vocab扩展") - elif o0 <= n0 and o1 <= n1: - # 权重扩展（old 是 new 的子集，放在左上角） - new_w[:o0, :o1] = old_w - if o0 < n0 or o1 < n1: - print(f" 扩展 {key}: ({o0},{o1})→({n0},{n1})") - elif n0 <= o0 and n1 <= o1: - # 权重缩小（new 是 old 的子集，取左上角） - new_w = old_w[:n0, :n1].clone() - print(f" 缩小 {key}: ({o0},{o1})→({n0},{n1})") - else: - # 特殊形状 → 尽量复制重合部分 - min0 = min(o0, n0) - min1 = min(o1, n1) - new_w[:min0, :min1] = old_w[:min0, :min1] - print(f" 部分迁移 {key}: ({o0},{o1})→({n0},{n1})") - # 3. 一维权重（LayerNorm, bias等） - elif old_w.dim() == 1: - o0 = old_w.shape[0] - n0 = new_w.shape[0] - min_d = min(o0, n0) - new_w[:min_d] = old_w[:min_d] - if o0 != n0: - print(f" 迁移 {key}: ({o0})→({n0}) norm扩展") - else: - new_w = old_w.clone() - else: - # 旧模型中没有这个key（可能是不存在的层） - # 检查是否是额外的新层 - # 如 transformer_blocks.6.xxx → 旧模型只有0-5 - new_sd[key] = new_w # 保持随机初始化 - continue - - new_sd[key] = new_w - - return new_sd - - -def main(): - print("=" * 60) - print("小狸训练 - 从头训练版 v2 (13.2M架构)") - print("=" * 60) - - device = 'cpu' - model_dir = '/workspace/xiaoli-llm-chat' - - # 加载数据 - 从合并文件加载 - merged_path = os.path.join(model_dir, 'xiaoli_merged_data_v2.json') - if not os.path.exists(merged_path): - print("❌ 未找到合并数据文件 xiaoli_merged_data.json") - print("请先运行合并脚本") - return - - with open(merged_path, 'r', encoding='utf-8') as f: - data = json.load(f) - - unique = dict_to_tuple(data) - print(f"加载合并数据: {len(unique)}条") - - # ====== 划分训练/验证集（保存到独立文件，避免问题重复问题 ====== - val_split = 0.1 - train_path = os.path.join(model_dir, 'xiaoli_train_v2.json') - val_path = os.path.join(model_dir, 'xiaoli_valid_v2.json') - - if os.path.exists(train_path) and os.path.exists(val_path): - # 从独立文件加载（精确的position-based划分） - with open(train_path, 'r', encoding='utf-8') as f: - train_set = dict_to_tuple(json.load(f)) - with open(val_path, 'r', encoding='utf-8') as f: - val_set = dict_to_tuple(json.load(f)) - # 检查是否有新数据（来自外部脚本） - train_qs = set(q for q, _ in train_set) - val_qs = set(q for q, _ in val_set) - new_entries = [(q, a) for q, a in unique if q not in train_qs and q not in val_qs] - if new_entries: - train_set.extend(new_entries) - with open(train_path, 'w', encoding='utf-8') as f: - json.dump([{"question": q, "answer": a} for q, a in train_set], f, ensure_ascii=False, indent=2) - print(f" 新增外部数据: {len(new_entries)}条 → 加入训练集") - print(f" 训练集: {len(train_set)}条 | 验证集: {len(val_set)}条 (来自文件)") - else: - # 首次划分：随机打乱后按位置切分 - shuffled = unique.copy() - random.seed(42) - random.shuffle(shuffled) - split_idx = int(len(shuffled) * val_split) - val_set = shuffled[:split_idx] - train_set = shuffled[split_idx:] - # 保存两份独立文件 - with open(train_path, 'w', encoding='utf-8') as f: - json.dump([{"question": q, "answer": a} for q, a in train_set], f, ensure_ascii=False, indent=2) - with open(val_path, 'w', encoding='utf-8') as f: - json.dump([{"question": q, "answer": a} for q, a in val_set], f, ensure_ascii=False, indent=2) - print(f" 首次划分: 训练集{len(train_set)}条 + 验证集{len(val_set)}条 (已保存)") - - print(f" 训练集: {len(train_set)}条 | 验证集: {len(val_set)}条") - - # ====== 加载固定vocab ====== - vocab_path = os.path.join(model_dir, 'vocab.json') - if not os.path.exists(vocab_path): - print("❌ 未找到 vocab.json，请先运行BPE词表生成脚本") - return - with open(vocab_path, 'r', encoding='utf-8') as f: - vocab = json.load(f) - print(f"词汇表: {len(vocab)} (BPE, 来自 vocab.json)") - # 同步tokenizer（新字符映射为UNK） - tokenizer_instance.char_to_id = vocab.copy() - tokenizer_instance.id_to_char = {v: k for k, v in vocab.items()} - tokenizer_instance._sorted_vocab_cache = None - for attr in ['_multi_trie', '_multi_keys', '_single_map']: - if hasattr(tokenizer_instance, attr): - delattr(tokenizer_instance, attr) - - # ====== 加载蒸馏问题(v3 + 真题) ====== - distill_q_path = os.path.join(model_dir, 'exam_distill_questions.json') - if not os.path.exists(distill_q_path): - distill_q_path = os.path.join(model_dir, 'distill_questions_v3.json') - if not os.path.exists(distill_q_path): - distill_q_path = os.path.join(model_dir, 'distill_questions_v2.json') - if os.path.exists(distill_q_path): - with open(distill_q_path, 'r', encoding='utf-8') as f: - distill_questions = json.load(f) - print(f"蒸馏问题: {len(distill_questions)}道") - else: - distill_questions = [] - print("蒸馏问题: 无") - - config = { - 'batch_size': 4, - 'gradient_accumulation_steps': 2, - 'epochs': 200, - 'd_model': 320, 'n_heads': 8, 'n_layers': 6, - 'd_ff': 1536, 'max_length': 256, 'dropout': 0.1, - 'use_amp': False, - } - - model = LLMModel( - vocab_size=len(vocab), d_model=config['d_model'], - n_heads=config['n_heads'], n_layers=config['n_layers'], - d_ff=config['d_ff'], max_len=config['max_length'], - dropout=config['dropout'], - use_memory=False, use_thinking=False, use_moe=False, - num_experts=1, top_k=1, - ) - total_params = sum(p.numel() for p in model.parameters()) - print(f"参数量: {total_params:,} ({total_params/1e6:.2f}M)") - - # 训练集/验证集 DataLoader - train_dataset = ConversationDataset(train_set, vocab, config['max_length']) - train_loader = DataLoader(train_dataset, batch_size=config['batch_size'], - shuffle=True, num_workers=0) - print(f"训练批次: {len(train_loader)}") - - val_dataset = ConversationDataset(val_set, vocab, config['max_length']) - val_loader = DataLoader(val_dataset, batch_size=config['batch_size'], - shuffle=False, num_workers=0) - print(f"验证批次: {len(val_loader)}") - - # 检查断点 - ckpt_path, start_epoch = find_latest_checkpoint(model_dir) - if ckpt_path and start_epoch > 0: - print(f"\n发现断点: epoch {start_epoch} ({ckpt_path})") - checkpoint = torch.load(ckpt_path, map_location=device) - old_cfg = checkpoint.get('model_config', {}) - old_d_model = old_cfg.get('d_model', 256) - old_n_layers = old_cfg.get('n_layers', 6) - old_d_ff = old_cfg.get('d_ff', 1024) - - # 检测架构变化，自动做权重扩展迁移 - new_sd_init = model.state_dict() # 随机初始化的新权重 - old_sd = checkpoint['model_state_dict'] - old_vocab_size = old_cfg.get('vocab_size', 0) - new_vocab_size = len(vocab) - - # 如果架构没变化，直接加载 - if (old_d_model == config['d_model'] and old_n_layers == config['n_layers'] - and old_d_ff == config['d_ff'] and old_vocab_size == new_vocab_size): - model.load_state_dict(old_sd) - print("✅ 架构一致，直接加载权重") - else: - print(f" 架构变化: d_model={old_d_model}→{config['d_model']}, " - f"layers={old_n_layers}→{config['n_layers']}, " - f"d_ff={old_d_ff}→{config['d_ff']}, " - f"vocab={old_vocab_size}→{new_vocab_size}") - print(" 执行权重扩展迁移（宽度扩展 + 深度复制）...") - new_sd = migrate_weights(old_sd, new_sd_init, old_d_model, config['d_model'], - old_n_layers, config['n_layers'], - old_d_ff, config['d_ff'], - old_vocab_size, new_vocab_size) - model.load_state_dict(new_sd) - print(" ✅ 权重扩展迁移完成！") - - # 记录旧checkpoint已使用，从下一epoch开始 - start_epoch = start_epoch - - # 尝试加载优化器（仅当架构vocab完全一致时） - if (old_d_model == config['d_model'] and old_n_layers == config['n_layers'] - and old_d_ff == config['d_ff'] and old_vocab_size == new_vocab_size - and 'optimizer_state_dict' in checkpoint): - try: - trainer = LLMTrainer(model, device, use_amp=config['use_amp']) - # 前向一步恢复优化器（不做完整epoch，太慢） - batch = next(iter(train_loader)) - input_ids = batch['input_ids'].to(device) - target_ids = batch['target_ids'].to(device) - mask = model.generate_mask(input_ids.size(1)).to(device) - outputs = model(input_ids, mask) - loss = trainer.criterion(outputs.reshape(-1, outputs.size(-1)), target_ids.reshape(-1)) - loss.backward() - trainer.optimizer.step() - trainer.optimizer.zero_grad() - # 加载优化器状态 - trainer.optimizer.load_state_dict(checkpoint['optimizer_state_dict']) - print("优化器状态已加载（1步热身）") - except Exception as e: - print(f"优化器加载失败（将使用新的）: {e}") - - print(f"\n开始训练: epoch {start_epoch + 1} -> {config['epochs']}") - print("=" * 60) - start = time.time() - - # 确保trainer已创建（可能因为架构变化没有在提前创建） - if 'trainer' not in dir() or trainer is None: - trainer = LLMTrainer(model, device, use_amp=config['use_amp']) - - best_val_loss = float('inf') - total_distill_added = 0 - for epoch in range(start_epoch + 1, config['epochs'] + 1): - # ====== 每轮从train_path重建dataloader（蒸馏数据已追加至此）====== - with open(train_path, 'r', encoding='utf-8') as f: - epoch_train = dict_to_tuple(json.load(f)) - print(f"\n 训练数据: {len(epoch_train)}条 (含{total_distill_added}条蒸馏新增)") - - train_dataset = ConversationDataset(epoch_train, vocab, config['max_length']) - train_loader = DataLoader(train_dataset, batch_size=config['batch_size'], - shuffle=True, num_workers=0) - - # ====== 训练 ====== - t0 = time.time() - avg_loss = trainer.train_epoch(train_loader, epoch, vocab, config) - dt = time.time() - t0 - - # ====== 验证 ====== - t1 = time.time() - val_loss = trainer.evaluate(val_loader, config) - dt_val = time.time() - t1 - - # 打印train/val loss - marker = "" - if val_loss < best_val_loss: - best_val_loss = val_loss - marker = " ★" - best_path = os.path.join(model_dir, 'xiaoli_v2_best.pth') - trainer.save_model(best_path, vocab, epoch) - print(f"Epoch {epoch}/{config['epochs']} | Train Loss: {avg_loss:.4f} | Val Loss: {val_loss:.4f}{marker} | {dt:.0f}s+{dt_val:.0f}s") - - # 每轮保存断点 - resume_path = os.path.join(model_dir, f'llm_model_v2_resume_epoch_{epoch}.pth') - trainer.save_model(resume_path, vocab, epoch) - if epoch % 20 == 0: - save_path = os.path.join(model_dir, f'llm_model_v2_epoch_{epoch}.pth') - trainer.save_model(save_path, vocab, epoch) - print(f" 已保存正式checkpoint: epoch {epoch}") - resume_files = sorted(glob.glob(os.path.join(model_dir, 'llm_model_v2_resume_*.pth'))) - for old in resume_files[:-5]: - try: os.remove(old) - except: pass - - # ====== 每轮末尾：在线蒸馏 ====== - if distill_questions: - distill_t0 = time.time() - try: - print(f" 📡 在线蒸馏 ({len(distill_questions)}题 × 多种方式)...") - chat = LLMChat(model, vocab, device, config['max_length']) - all_new_pairs = [] - - # ====== Pass 1: 模型生成(temperature=0.15) → API转换 ====== - print(" [P1] 模型生成中...") - model_answers = [] - for i, q in enumerate(distill_questions): - ans = chat.generate_response(q, max_length=40, temperature=0.15) - model_answers.append((q, ans)) - if (i + 1) % 100 == 0: - print(f" P1: {i+1}/{len(distill_questions)}") - p1_gen = time.time() - print(f" [P1] 生成完毕 ({(p1_gen-distill_t0):.0f}s)，API转换中...") - - def progress(done, total): - if done % 50 == 0 or done == total: - print(f" API: {done}/{total}", end="\r") - - p1_results = batch_convert(model_answers, progress) - print(f" API: {len(model_answers)}/{len(model_answers)}") - all_new_pairs.extend(p1_results) - print(f" [P1] +{len(p1_results)}条") - - # ====== Pass 2: 模型高温生成(temperature=0.7) → 更多多样性 ====== - print(f" [P2] 模型高温生成中 (temp=0.7)...") - model_answers2 = [] - for i, q in enumerate(distill_questions): - ans = chat.generate_response(q, max_length=40, temperature=0.7) - model_answers2.append((q, ans)) - if (i + 1) % 100 == 0: - print(f" P2: {i+1}/{len(distill_questions)}") - p2_gen = time.time() - print(f" [P2] 生成完毕 ({(p2_gen-p1_gen):.0f}s)，API转换中...") - - p2_results = batch_convert(model_answers2, progress) - print(f" API: {len(model_answers2)}/{len(model_answers2)}") - all_new_pairs.extend(p2_results) - print(f" [P2] +{len(p2_results)}条") - - # ====== Pass 3: API直接生成多样回答(temp=1.0) ====== - print(f" [P3] API直接生成 (temp=1.0, 全部直答)...") - p3_t0 = time.time() - p3_results = batch_generate_multi_temp( - distill_questions, - temperatures=[1.0], - max_workers=8 - ) - all_new_pairs.extend(p3_results) - print(f" [P3] +{len(p3_results)}条") - - # ====== 合并去重后追加 ====== - if all_new_pairs: - with open(train_path, 'r', encoding='utf-8') as f: - cur_train = json.load(f) - existing_pairs = set((item.get("question",""), item.get("answer","")) for item in cur_train) - truly_new = [item for item in all_new_pairs - if (item["question"], item["answer"]) not in existing_pairs] - - if truly_new: - cur_train.extend(truly_new) - with open(train_path, 'w', encoding='utf-8') as f: - json.dump(cur_train, f, ensure_ascii=False, indent=2) - # 同步到merged - with open(merged_path, 'r', encoding='utf-8') as f: - cur_merged = json.load(f) - cur_merged.extend(truly_new) - with open(merged_path, 'w', encoding='utf-8') as f: - json.dump(cur_merged, f, ensure_ascii=False, indent=2) - total_distill_added += len(truly_new) - distill_total = time.time() - distill_t0 - print(f" ✅ 蒸馏完成: +{len(truly_new)}条 (新增{len(set(q for q,a in [(i['question'],i['answer']) for i in truly_new]))}道独特问题) | 总耗时{distill_total:.0f}s") - else: - print(f" ℹ️ 蒸馏完成: 本轮无新内容（全部已存在）({time.time()-distill_t0:.0f}s)") - else: - print(f" ℹ️ 蒸馏完成: 无有效转换 ({time.time()-distill_t0:.0f}s)") - except Exception as e: - print(f" ⚠️ 蒸馏异常: {e}") - - # 保存最终模型 - trainer.save_model(os.path.join(model_dir, 'xiaoli_v2_final.pth'), vocab, config['epochs']) - - # 清理所有resume文件 - for f in glob.glob(os.path.join(model_dir, 'llm_model_v2_resume_*.pth')): - try: - os.remove(f) - except: - pass - - total = time.time() - start - print(f"\n训练完成！总耗时: {total/60:.1f}分钟") - - # 测试 - print("\n测试:") - chat = LLMChat(model, vocab, device, config['max_length']) - for q in ["你好", "你是谁", "你叫什么名字", "你好可爱", "我喜欢你", "1+1等于几"]: - r = chat.generate_response(q, max_length=60, temperature=0.8) - print(f" Q: {q}") - print(f" A: {r}\n") - -if __name__ == "__main__": - main() diff --git a/train_cloud.py b/train_cloud.py deleted file mode 100644 index 81591cb..0000000 --- a/train_cloud.py +++ /dev/null @@ -1,437 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -XiaLi 云端 GPU 训练脚本 v1.0 -============================= -两阶段训练 + ModelScope 自动下载数据 + 断点续训 - -用法（云GPU上）: - python3 train_cloud.py # 默认训练 - python3 train_cloud.py --resume # 从最新检查点续训 - python3 train_cloud.py --epochs 10 # 指定预训练轮数 -""" -import sys, os, json, time, glob, re, math, random -from pathlib import Path - -import torch -import torch.nn as nn -from torch.utils.data import Dataset, DataLoader -import numpy as np - -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -from llm_chat_improved import ( - CharTokenizer, LLMModel, LLMChat, - ConversationDataset, tokenizer_instance -) - -# ── 配置 ────────────────────────────────────────────────── -class CloudConfig: - # 模型架构 - d_model: int = 320 - n_heads: int = 8 - n_layers: int = 6 - d_ff: int = 1536 - max_length: int = 256 # 对齐用 - dropout: float = 0.1 - use_memory: bool = False # 关掉记忆/思考/MoE 加速 - use_thinking: bool = False - use_moe: bool = False - - # 预训练 - pt_batch_size: int = 64 # 32G显存可跑64, 24G改32 - pt_grad_accum: int = 1 # 大batch不需要累积 - pt_epochs: int = 5 - pt_lr: float = 5e-4 - pt_min_lr: float = 5e-5 - pt_warmup_steps: int = 200 - pt_max_length: int = 256 # 长序列学更好 - pt_weight_decay: float = 0.1 - pt_max_norm: float = 1.0 - - # 对齐 - align_batch_size: int = 32 - align_grad_accum: int = 1 - align_epochs: int = 20 - align_lr: float = 1e-4 - align_min_lr: float = 1e-5 - align_warmup_steps: int = 50 - align_weight_decay: float = 0.01 - align_max_norm: float = 1.0 - - # 通用 - device: str = "cuda" if torch.cuda.is_available() else "cpu" - num_workers: int = 4 - use_amp: bool = True - seed: int = 42 - save_dir: str = "/root/xiaoli-llm-chat" # AutoDL 默认路径 - log_every: int = 10 - save_interval: int = 1 - keep_checkpoints: int = 3 - - -# ── 词表加载 ────────────────────────────────────────────── -def load_vocab(vocab_path: str) -> dict: - special = {'': 0, '': 1, '': 2, '': 3, - '': 4, '': 5, '': 6} - with open(vocab_path, 'r') as f: - normal_vocab = json.load(f) - full_vocab = {**special, **normal_vocab} - # 同步到全局 tokenizer - tokenizer_instance.char_to_id = full_vocab.copy() - tokenizer_instance.id_to_char = {v: k for k, v in full_vocab.items()} - for attr in ['_multi_trie', '_multi_keys', '_single_map']: - if hasattr(tokenizer_instance, attr): delattr(tokenizer_instance, attr) - return full_vocab - - -# ── 预训练数据集（流式加载，不爆内存）──────────────────── -class PretrainDataset(Dataset): - def __init__(self, data_dir: str, vocab: dict, max_length: int = 128): - self.vocab = vocab - self.max_length = max_length - self.pad_id = vocab.get('', 0) - self.texts = [] - data_dir = Path(data_dir) - for fpath in sorted(data_dir.glob("*")): - if fpath.suffix in ('.txt', '.json', '.jsonl'): - self._load_file(fpath) - if not self.texts: - raise ValueError(f"预训练数据目录为空: {data_dir}") - total_chars = sum(len(t) for t in self.texts) - print(f" 预训练数据: {len(self.texts):,} 条, {total_chars:,} 字符") - - def _load_file(self, path: Path): - try: - if path.suffix == '.jsonl': - with open(path, 'r', encoding='utf-8', errors='ignore') as f: - for line in f: - try: - item = json.loads(line) - t = item.get('text', '').strip() - if len(t) >= 10: self.texts.append(t) - except: continue - except: pass - - def __len__(self): return len(self.texts) - - def __getitem__(self, idx): - text = self.texts[idx] - tokens = tokenizer_instance.encode(text) - if len(tokens) > self.max_length: - start = random.randint(0, len(tokens) - self.max_length) - tokens = tokens[start:start + self.max_length] - else: - tokens = tokens + [self.pad_id] * (self.max_length - len(tokens)) - input_ids = torch.tensor(tokens[:-1], dtype=torch.long) - target_ids = torch.tensor(tokens[1:], dtype=torch.long) - loss_mask = torch.ones_like(target_ids) - return {'input_ids': input_ids, 'target_ids': target_ids, - 'loss_mask': loss_mask, 'weights': 1.0} - - -# ── 训练器 ────────────────────────────────────────────── -class CloudTrainer: - def __init__(self, model: nn.Module, config: CloudConfig, stage: str = "pretrain"): - self.model = model.to(config.device) - self.device = config.device - self.config = config - self.stage = stage - self.criterion = nn.CrossEntropyLoss(ignore_index=0, label_smoothing=0.1) - - if stage == "pretrain": - base_lr = config.pt_lr - wd = config.pt_weight_decay - else: - base_lr = config.align_lr - wd = config.align_weight_decay - - # 分层学习率：embedding/fc_out 用半速 - param_groups = [] - for name, param in model.named_parameters(): - if not param.requires_grad: continue - if 'embedding' in name or 'fc_out' in name: - param_groups.append({'params': param, 'lr': base_lr * 0.5, 'weight_decay': wd}) - elif 'norm' in name or 'rmsnorm' in name: - param_groups.append({'params': param, 'lr': base_lr, 'weight_decay': 0.0}) - else: - param_groups.append({'params': param, 'lr': base_lr, 'weight_decay': wd}) - - self.optimizer = torch.optim.AdamW(param_groups, lr=base_lr, betas=(0.9, 0.95), eps=1e-8) - self.scaler = torch.cuda.amp.GradScaler(enabled=True) - self.scheduler = None - - if torch.cuda.is_available(): - torch.backends.cuda.matmul.allow_tf32 = True - torch.backends.cudnn.allow_tf32 = True - torch.backends.cudnn.benchmark = True - - def build_scheduler(self, total_steps: int): - if self.stage == "pretrain": - warmup = self.config.pt_warmup_steps - min_lr = self.config.pt_min_lr - else: - warmup = self.config.align_warmup_steps - min_lr = self.config.align_min_lr - base_lr = self.optimizer.param_groups[0]['lr'] - - def lr_lambda(step): - if step < warmup: - return step / max(1, warmup) - progress = (step - warmup) / max(1, total_steps - warmup) - cosine = 0.5 * (1 + math.cos(math.pi * progress)) - return (1 - min_lr / base_lr) * cosine + min_lr / base_lr - - self.scheduler = torch.optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda) - - def train_epoch(self, dataloader, epoch, total_epochs, stage_name="pretrain"): - self.model.train() - total_loss = 0 - n_batches = len(dataloader) - accum_steps = self.config.pt_grad_accum if stage_name == "pretrain" else self.config.align_grad_accum - max_norm = self.config.pt_max_norm if stage_name == "pretrain" else self.config.align_max_norm - - self.optimizer.zero_grad() - for batch_idx, batch in enumerate(dataloader): - input_ids = batch['input_ids'].to(self.device, non_blocking=True) - target_ids = batch['target_ids'].to(self.device, non_blocking=True) - mask = self.model.generate_mask(input_ids.size(1)).to(self.device, non_blocking=True) - - with torch.amp.autocast('cuda', enabled=True): - outputs = self.model(input_ids, mask) - loss = self.criterion(outputs.reshape(-1, outputs.size(-1)), target_ids.reshape(-1)) - loss = loss / accum_steps - - self.scaler.scale(loss).backward() - - if (batch_idx + 1) % accum_steps == 0: - self.scaler.unscale_(self.optimizer) - torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm) - self.scaler.step(self.optimizer) - self.scaler.update() - self.optimizer.zero_grad() - if self.scheduler: self.scheduler.step() - - total_loss += loss.item() * accum_steps - - if (batch_idx + 1) % max(1, n_batches // 10) == 0: - lr = self.optimizer.param_groups[0]['lr'] - mem = f" | Mem: {torch.cuda.max_memory_allocated()/1024**3:.1f}GiB" if torch.cuda.is_available() else "" - print(f" [{stage_name}] Epoch {epoch}/{total_epochs} | " - f"Batch {batch_idx+1}/{n_batches} | Loss: {total_loss/(batch_idx+1):.4f} | LR: {lr:.2e}{mem}") - - return total_loss / n_batches - - def save_checkpoint(self, path: str, vocab: dict, epoch: int, loss: float): - core = getattr(self.model, '_orig_mod', self.model) - torch.save({ - 'model_state_dict': core.state_dict(), - 'optimizer_state_dict': self.optimizer.state_dict(), - 'scheduler_state_dict': self.scheduler.state_dict() if self.scheduler else None, - 'vocab': vocab, - 'epoch': epoch, - 'loss': loss, - 'config': {k: getattr(self.config, k) for k in dir(self.config) - if not k.startswith('_') and not callable(getattr(self.config, k))}, - }, path) - print(f" ✅ 已保存: {path} (epoch {epoch}, loss {loss:.4f})") - - def load_checkpoint(self, path: str) -> tuple: - cp = torch.load(path, map_location=self.device, weights_only=False) - core = getattr(self.model, '_orig_mod', self.model) - core.load_state_dict(cp['model_state_dict']) - self.optimizer.load_state_dict(cp['optimizer_state_dict']) - if cp.get('scheduler_state_dict') and self.scheduler: - self.scheduler.load_state_dict(cp['scheduler_state_dict']) - return cp.get('epoch', 0), cp.get('loss', float('inf')), cp.get('vocab', {}) - - -# ── 数据下载 ────────────────────────────────────────────── -def download_data(save_dir: str): - """从 ModelScope 下载预训练数据""" - data_dir = os.path.join(save_dir, 'pretrain_data') - if os.path.exists(data_dir) and list(Path(data_dir).glob("*.jsonl")): - print(" 预训练数据已存在，跳过下载") - return data_dir - - os.makedirs(data_dir, exist_ok=True) - print(" 📥 从 ModelScope 下载预训练数据...") - try: - import subprocess - subprocess.run([ - sys.executable, '-m', 'pip', 'install', 'modelscope', '-q' - ], check=True, capture_output=True) - subprocess.run([ - sys.executable, '-m', 'modelscope', 'download', - '--dataset', 'gongjy/minimind_dataset', - '--include', 'pretrain_t2t_mini*', - '--local_dir', data_dir, - ], check=True) - print(f" ✅ 下载完成: {data_dir}") - except Exception as e: - print(f" ⚠️ 自动下载失败: {e}") - print(f" ⚠️ 请手动下载放到 {data_dir}") - print(f" ⚠️ 命令: modelscope download --dataset gongjy/minimind_dataset --include 'pretrain_t2t_mini*' --local_dir {data_dir}") - return data_dir - - -# ── 主函数 ────────────────────────────────────────────── -def main(): - import argparse - parser = argparse.ArgumentParser() - parser.add_argument("--resume", action="store_true", help="从最新检查点续训") - parser.add_argument("--epochs", type=int, default=5, help="预训练轮数") - parser.add_argument("--align-epochs", type=int, default=20, help="对齐轮数") - parser.add_argument("--save-dir", type=str, default="/root/xiaoli-llm-chat") - parser.add_argument("--batch-size", type=int, default=0, help="覆盖batch大小") - args = parser.parse_args() - - config = CloudConfig() - config.pt_epochs = args.epochs - config.align_epochs = args.align_epochs - config.save_dir = args.save_dir - - if args.batch_size > 0: - config.pt_batch_size = args.batch_size - config.align_batch_size = args.batch_size // 2 - - if not torch.cuda.is_available(): - print("⚠️ 未检测到 GPU！云端训练建议使用 AutoDL / 矩池云 / 恒源云") - - print("=" * 65) - print(f" XiaLi 云端训练") - print(f" GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'N/A'}") - print(f" 显存: {torch.cuda.get_device_properties(0).total_mem / 1024**3:.1f} GiB" if torch.cuda.is_available() else "") - print(f" 模型: {config.d_model}d, {config.n_layers}L, {config.n_heads}H") - print(f" 预训练: {config.pt_epochs} epochs, batch={config.pt_batch_size}") - print(f" 对齐: {config.align_epochs} epochs, batch={config.align_batch_size}") - print("=" * 65) - - os.makedirs(config.save_dir, exist_ok=True) - - # 1. 加载词表 - vocab_path = os.path.join(config.save_dir, 'vocab.json') - if not os.path.exists(vocab_path): - print(f"❌ 未找到 {vocab_path}") - print(f" 请从 Gitee 仓库获取: https://gitee.com/shuiyu1123/xiaoli-llm-chat") - return - vocab = load_vocab(vocab_path) - print(f"词表: {len(vocab)} tokens, vocab_size={max(vocab.values())+1}") - - # 2. 创建模型 - model = LLMModel( - vocab_size=max(vocab.values())+1, - d_model=config.d_model, - n_heads=config.n_heads, - n_layers=config.n_layers, - d_ff=config.d_ff, - max_len=config.max_length, - dropout=config.dropout, - use_memory=False, use_thinking=False, use_moe=False, - ) - params = sum(p.numel() for p in model.parameters()) - print(f"参数量: {params:,} ({params/1e6:.2f}M)") - - # 3. 续训检测 - start_epoch = 1 - if args.resume: - checkpoints = sorted(glob.glob(os.path.join(config.save_dir, 'pretrain_epoch_*.pth'))) - if checkpoints: - latest = checkpoints[-1] - trainer_pt = CloudTrainer(model, config, stage="pretrain") - loaded_epoch, loaded_loss, _ = trainer_pt.load_checkpoint(latest) - start_epoch = loaded_epoch + 1 - print(f" 从 {latest} 续训 (epoch {loaded_epoch}, loss {loaded_loss:.4f})") - - # 4. 下载数据 - pt_dir = download_data(config.save_dir) - - # 5. 预训练 - print(f"\n{'='*65}") - print(f" 阶段一：预训练 ({config.pt_epochs} epochs)") - print(f"{'='*65}") - - pt_dataset = PretrainDataset(pt_dir, vocab, config.pt_max_length) - pt_loader = DataLoader( - pt_dataset, batch_size=config.pt_batch_size, shuffle=True, - num_workers=config.num_workers, - pin_memory=True, - ) - - trainer = CloudTrainer(model, config, stage="pretrain") - total_steps = len(pt_loader) * config.pt_epochs - trainer.build_scheduler(total_steps) - - for epoch in range(start_epoch, config.pt_epochs + 1): - t0 = time.time() - loss = trainer.train_epoch(pt_loader, epoch, config.pt_epochs, "pretrain") - epoch_time = time.time() - t0 - lr = trainer.optimizer.param_groups[0]['lr'] - print(f" → Epoch {epoch}/{config.pt_epochs} | Loss: {loss:.4f} | LR: {lr:.2e} | 耗时: {epoch_time:.0f}s") - - if epoch % config.save_interval == 0: - ckpt_path = os.path.join(config.save_dir, f'pretrain_epoch_{epoch}.pth') - trainer.save_checkpoint(ckpt_path, vocab, epoch, loss) - # 清理旧检查点 - all_ckpts = sorted(glob.glob(os.path.join(config.save_dir, 'pretrain_epoch_*.pth'))) - for old in all_ckpts[:-config.keep_checkpoints]: - os.remove(old) - - print(f"\n ✅ 预训练完成") - - # 6. 对齐 - print(f"\n{'='*65}") - print(f" 阶段二：角色对齐 ({config.align_epochs} epochs)") - print(f"{'='*65}") - - merged_path = os.path.join(config.save_dir, 'xiaoli_merged_data_v2.json') - if not os.path.exists(merged_path): - print(f" ⚠️ 未找到合并数据: {merged_path}") - print(f" ⚠️ 跳过对齐阶段") - else: - with open(merged_path, 'r') as f: - raw_data = json.load(f) - pairs = [] - for item in raw_data: - q = item.get("question", "") or item.get("input", "") or item.get("prompt", "") - a = item.get("answer", "") or item.get("output", "") or item.get("response", "") - if q and a: pairs.append((q, a)) - - random.seed(config.seed) - random.shuffle(pairs) - split = int(len(pairs) * 0.95) - train_pairs = pairs[:split] - print(f" 对齐数据: {len(pairs)} 条 (训练 {len(train_pairs)} + 验证 {len(pairs)-split})") - - train_ds = ConversationDataset(train_pairs, vocab, config.max_length) - train_loader = DataLoader( - train_ds, batch_size=config.align_batch_size, shuffle=True, - num_workers=config.num_workers, pin_memory=True, - ) - - trainer2 = CloudTrainer(model, config, stage="align") - total_align_steps = len(train_loader) * config.align_epochs - trainer2.build_scheduler(total_align_steps) - - for epoch in range(1, config.align_epochs + 1): - t0 = time.time() - loss = trainer2.train_epoch(train_loader, epoch, config.align_epochs, "align") - epoch_time = time.time() - t0 - print(f" → Epoch {epoch}/{config.align_epochs} | Loss: {loss:.4f} | 耗时: {epoch_time:.0f}s") - - if epoch % 5 == 0: - ckpt_path = os.path.join(config.save_dir, f'xiaoli_align_epoch_{epoch}.pth') - trainer2.save_checkpoint(ckpt_path, vocab, epoch, loss) - - # 保存最终模型 - final_path = os.path.join(config.save_dir, 'xiaoli_final.pth') - trainer2.save_checkpoint(final_path, vocab, config.align_epochs, loss) - print(f"\n ✅ 对齐完成！最终模型: {final_path}") - - print(f"\n{'='*65}") - print(f" 🎉 训练完成！") - print(f" 模型保存在: {config.save_dir}") - print(f" {'='*65}") - - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/train_max.py b/train_max.py deleted file mode 100644 index 473ce52..0000000 --- a/train_max.py +++ /dev/null @@ -1,662 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -小狸 AI - 最大化训练脚本 v2.0 -============================== -两阶段训练：预训练(通用中文) → 角色对齐(QA数据) - -相比原始 train.py 的改进： - ① 两阶段训练：先用通用语料建立语言能力，再用QA数据做角色对齐 - ② Warmup + Cosine 学习率调度（业界标准） - ③ 更大的有效 batch size（提升收敛稳定性） - ④ 启用 AMP 混合精度（GPU加速） - ⑤ DataLoader 多进程加载 - ⑥ 修复在线蒸馏 P2/P3 失效 bug - ⑦ 支持学习率分层（不同层不同 LR） - ⑧ 更激进的梯度裁剪策略 - ⑨ 支持从原始 train.py checkpoint 无缝续训 -""" - -import sys, os, json, time, glob, random, re, math -from typing import List, Tuple, Optional -from pathlib import Path - -import torch -import torch.nn as nn -from torch.utils.data import Dataset, DataLoader -import numpy as np - -# 复用原始项目代码 -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -from llm_chat_improved import ( - CharTokenizer, LLMModel, LLMChat, - ConversationDataset, tokenizer_instance -) -from api_correction import call_astron - -# ── 配置 ────────────────────────────────────────────────── -class Config: - # ========= 模型架构 ========= - d_model: int = 320 - n_heads: int = 8 - n_layers: int = 6 - d_ff: int = 1536 - max_length: int = 256 # 训练序列长度 - dropout: float = 0.1 - use_memory: bool = False - use_thinking: bool = False - use_moe: bool = False - num_experts: int = 1 - top_k: int = 1 - - # ========= 阶段一：预训练（通用中文语料） ========= - pt_enable: bool = True # 是否启用预训练阶段 - pt_data_dir: str = "" # 预训练数据目录（留空则自动下载） - pt_batch_size: int = 16 # 预训练 batch - pt_grad_accum: int = 4 # 梯度累积 - pt_effective_batch: int = 64 # = pt_batch_size * pt_grad_accum - pt_epochs: int = 5 # 预训练 epoch - pt_lr: float = 5e-4 # 预训练初始 LR（比微调大） - pt_min_lr: float = 5e-5 # 预训练最小 LR - pt_warmup_steps: int = 500 # 预训练 warmup - pt_weight_decay: float = 0.1 # 预训练 weight decay（标准） - pt_max_norm: float = 1.0 # 梯度裁剪 - pt_max_length: int = 128 # 预训练用较短序列（加速） - pt_data_ratio: float = 1.0 # 使用数据的比例（调试用） - - # ========= 阶段二：角色对齐（QA数据） ========= - align_batch_size: int = 8 - align_grad_accum: int = 4 - align_effective_batch: int = 32 - align_epochs: int = 50 - align_lr: float = 1e-4 - align_min_lr: float = 1e-5 - align_warmup_steps: int = 100 - align_weight_decay: float = 0.01 - align_max_norm: float = 1.0 - - # ========= 在线蒸馏（阶段二末尾执行） ========= - distill_enable: bool = True - distill_every: int = 5 # 每 N 个 epoch 执行一次蒸馏 - - # ========= 通用 ========= - device: str = "cuda" if torch.cuda.is_available() else "cpu" - num_workers: int = 4 if torch.cuda.is_available() else 0 - use_amp: bool = torch.cuda.is_available() - seed: int = 42 - save_dir: str = os.path.dirname(os.path.abspath(__file__)) - log_every: int = 0 # 0 = 每1%打印 - - # ========= 检查点 ========= - save_interval: int = 1 # 每 N 个 epoch 保存 resume - archive_interval: int = 10 # 每 N 个 epoch 保存归档 - keep_resume: int = 3 # 保留最近 resume 数 - - -config = Config() - - -# ── 预训练数据集 ────────────────────────────────────────── -class PretrainDataset(Dataset): - """通用中文文本预训练数据集 - - 支持两种输入格式： - 1. 纯文本文件：每行一句话/段落 - 2. JSON 文件：列表，每个元素含 "text" 字段 - """ - - def __init__(self, data_dir: str, vocab: dict, max_length: int = 128): - self.vocab = vocab - self.max_length = max_length - self.pad_id = vocab.get('', 0) - - # 扫描并加载所有文本 - self.texts = [] - data_dir = Path(data_dir) - if not data_dir.exists(): - raise FileNotFoundError(f"预训练数据目录不存在: {data_dir}") - - for fpath in sorted(data_dir.glob("*")): - if fpath.suffix in ('.txt', '.json', '.jsonl'): - self._load_file(fpath) - - if not self.texts: - raise ValueError(f"预训练数据目录为空: {data_dir}") - - print(f"预训练数据: {len(self.texts)} 条文本, " - f"总字符数: {sum(len(t) for t in self.texts):,}") - - def _load_file(self, path: Path): - try: - if path.suffix == '.txt': - with open(path, 'r', encoding='utf-8', errors='ignore') as f: - for line in f: - line = line.strip() - if len(line) >= 20: # 过滤太短的 - self.texts.append(line) - elif path.suffix == '.json': - with open(path, 'r', encoding='utf-8', errors='ignore') as f: - data = json.load(f) - for item in data: - if isinstance(item, dict) and 'text' in item: - t = item['text'].strip() - if len(t) >= 20: - self.texts.append(t) - elif isinstance(item, str) and len(item) >= 20: - self.texts.append(item) - elif path.suffix == '.jsonl': - with open(path, 'r', encoding='utf-8', errors='ignore') as f: - for line in f: - try: - item = json.loads(line) - if isinstance(item, dict) and 'text' in item: - t = item['text'].strip() - if len(t) >= 20: - self.texts.append(t) - except json.JSONDecodeError: - continue - except Exception as e: - print(f" 跳过 {path.name}: {e}") - - def __len__(self): - return len(self.texts) - - def __getitem__(self, idx): - text = self.texts[idx] - # 随机截取一段（提升多样性） - tokens = tokenizer_instance.encode(text) - if len(tokens) > self.max_length: - start = random.randint(0, len(tokens) - self.max_length) - tokens = tokens[start:start + self.max_length] - else: - tokens = tokens + [self.pad_id] * (self.max_length - len(tokens)) - - # 自回归：input = tokens[:-1], target = tokens[1:] - input_ids = torch.tensor(tokens[:-1], dtype=torch.long) - target_ids = torch.tensor(tokens[1:], dtype=torch.long) - return {'input_ids': input_ids, 'target_ids': target_ids} - - -# ── 改进版训练器 ────────────────────────────────────────── -class MaxTrainer: - """增强训练器 — 支持 warmup、分层 LR、自定义 scheduler""" - - def __init__(self, model: nn.Module, config: Config, stage: str = "pretrain"): - self.model = model.to(config.device) - self.device = config.device - self.config = config - self.stage = stage - self.criterion = nn.CrossEntropyLoss(ignore_index=0, label_smoothing=0.1) - - # 根据阶段选择超参数 - if stage == "pretrain": - base_lr = config.pt_lr - wd = config.pt_weight_decay - else: - base_lr = config.align_lr - wd = config.align_weight_decay - - # 分层学习率：embedding + fc_out 用低 LR，中间层用高 LR - param_groups = [] - for name, param in model.named_parameters(): - if not param.requires_grad: - continue - if 'embedding' in name or 'fc_out' in name: - param_groups.append({'params': param, 'lr': base_lr * 0.5, 'weight_decay': wd}) - elif 'norm' in name or 'rmsnorm' in name: - param_groups.append({'params': param, 'lr': base_lr, 'weight_decay': 0.0}) - else: - param_groups.append({'params': param, 'lr': base_lr, 'weight_decay': wd}) - - self.optimizer = torch.optim.AdamW(param_groups, lr=base_lr, weight_decay=wd, - betas=(0.9, 0.95), eps=1e-8) - - # AMP - self.use_amp = config.use_amp and torch.cuda.is_available() - self.scaler = torch.cuda.amp.GradScaler(enabled=self.use_amp) - - if torch.cuda.is_available(): - torch.backends.cuda.matmul.allow_tf32 = True - torch.backends.cudnn.allow_tf32 = True - torch.backends.cudnn.benchmark = True - torch.backends.cudnn.deterministic = False - - def build_scheduler(self, total_steps: int): - """构建 Warmup + Cosine Decay scheduler""" - if self.stage == "pretrain": - warmup = self.config.pt_warmup_steps - min_lr = self.config.pt_min_lr - else: - warmup = self.config.align_warmup_steps - min_lr = self.config.align_min_lr - - def lr_lambda(step): - if step < warmup: - return step / max(1, warmup) - progress = (step - warmup) / max(1, total_steps - warmup) - cosine_decay = 0.5 * (1 + math.cos(math.pi * progress)) - return (1 - min_lr / self.optimizer.param_groups[0]['lr']) * cosine_decay \ - + min_lr / self.optimizer.param_groups[0]['lr'] - - self.scheduler = torch.optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda) - - def train_epoch(self, dataloader: DataLoader, epoch: int, total_epochs: int, - stage_name: str = "pretrain", log_every: int = 0) -> float: - self.model.train() - total_loss = 0 - n_batches = len(dataloader) - - if stage_name == "pretrain": - accum = self.config.pt_grad_accum - max_norm = self.config.pt_max_norm - else: - accum = self.config.align_grad_accum - max_norm = self.config.align_max_norm - - if log_every == 0: - log_every = max(1, n_batches // 100) - - self.optimizer.zero_grad() - - for batch_idx, batch in enumerate(dataloader): - input_ids = batch['input_ids'].to(self.device, non_blocking=True) - target_ids = batch['target_ids'].to(self.device, non_blocking=True) - - mask = self.model.generate_mask(input_ids.size(1)).to(self.device) - - with torch.cuda.amp.autocast(enabled=self.use_amp): - outputs = self.model(input_ids, mask) - outputs = outputs.reshape(-1, outputs.size(-1)) - tgt = target_ids.reshape(-1) - loss = self.criterion(outputs, tgt) / accum - - self.scaler.scale(loss).backward() - - if (batch_idx + 1) % accum == 0: - self.scaler.unscale_(self.optimizer) - torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm) - self.scaler.step(self.optimizer) - self.scaler.update() - self.optimizer.zero_grad() - if self.scheduler is not None: - self.scheduler.step() - - total_loss += loss.item() * accum - - if (batch_idx + 1) % log_every == 0: - avg = total_loss / (batch_idx + 1) - lr = self.optimizer.param_groups[0]['lr'] - mem = "" - if torch.cuda.is_available(): - mem = f" | Mem: {torch.cuda.memory_allocated()/1e9:.2f}GB" - print(f" [{stage_name}] Epoch {epoch}/{total_epochs} | " - f"Batch {batch_idx+1}/{n_batches} | Loss: {avg:.4f} | LR: {lr:.2e}{mem}") - - return total_loss / n_batches - - def evaluate(self, dataloader: DataLoader) -> float: - self.model.eval() - total_loss = 0 - with torch.no_grad(): - for batch in dataloader: - input_ids = batch['input_ids'].to(self.device) - target_ids = batch['target_ids'].to(self.device) - mask = self.model.generate_mask(input_ids.size(1)).to(self.device) - outputs = self.model(input_ids, mask) - outputs = outputs.reshape(-1, outputs.size(-1)) - loss = self.criterion(outputs, target_ids.reshape(-1)) - total_loss += loss.item() - return total_loss / max(1, len(dataloader)) - - def save_model(self, filepath: str, vocab: dict, epoch: int): - core = getattr(self.model, '_orig_mod', self.model) - torch.save({ - 'model_state_dict': core.state_dict(), - 'optimizer_state_dict': self.optimizer.state_dict(), - 'scheduler_state_dict': self.scheduler.state_dict() if self.scheduler else None, - 'vocab': vocab, - 'model_config': { - 'vocab_size': core.vocab_size, - 'd_model': core.d_model, - 'n_heads': core.n_heads, - 'n_layers': core.n_layers, - 'd_ff': core.d_ff, - 'max_len': core.max_len, - 'use_memory': core.use_memory, - 'use_thinking': core.use_thinking, - 'use_moe': core.use_moe, - 'stage': self.stage, - }, - 'epoch': epoch, - }, filepath) - print(f" ✅ 已保存: {filepath}") - - -# ── 在线蒸馏（修复版） ────────────────────────────────── -def run_distillation(model, vocab, device, max_length, distill_questions, save_dir): - if not distill_questions: - return [], 0 - - from api_correction import batch_convert - chat = LLMChat(model, vocab, device, max_length) - all_new = [] - - print(f" 📡 蒸馏 {len(distill_questions)} 题...") - answers1 = [] - for i, q in enumerate(distill_questions): - ans = chat.generate_response(q, max_length=40, temperature=0.15) - answers1.append((q, ans)) - if (i + 1) % 100 == 0: - print(f" P1: {i+1}/{len(distill_questions)}") - - results1 = batch_convert(answers1) - all_new.extend(results1) - print(f" P1: +{len(results1)} 条") - - print(f" [P2] 模型高温生成 (temp=0.8)...") - for i, q in enumerate(distill_questions): - ans = chat.generate_response(q, max_length=40, temperature=0.8) - if ans and len(ans) >= 5: - all_new.append({"question": q, "answer": ans}) - if (i + 1) % 100 == 0: - print(f" P2: {i+1}/{len(distill_questions)}") - - print(f" [P3] API 直答 (temp=1.0)...") - sampled = random.sample(distill_questions, min(50, len(distill_questions))) - for q in sampled: - try: - ans = call_astron(q, temperature=1.0) - if ans and len(ans) >= 5: - all_new.append({"question": q, "answer": ans}) - except Exception: - pass - - train_path = os.path.join(save_dir, 'xiaoli_train_v2.json') - merged_path = os.path.join(save_dir, 'xiaoli_merged_data_v2.json') - added = 0 - - if all_new and os.path.exists(train_path): - with open(train_path, 'r', encoding='utf-8') as f: - cur_train = json.load(f) - existing = set((item.get("question", ""), item.get("answer", "")) for item in cur_train) - truly_new = [item for item in all_new - if (item["question"], item["answer"]) not in existing] - - if truly_new: - cur_train.extend(truly_new) - with open(train_path, 'w', encoding='utf-8') as f: - json.dump(cur_train, f, ensure_ascii=False, indent=2) - if os.path.exists(merged_path): - with open(merged_path, 'r', encoding='utf-8') as f: - cur_merged = json.load(f) - cur_merged.extend(truly_new) - with open(merged_path, 'w', encoding='utf-8') as f: - json.dump(cur_merged, f, ensure_ascii=False, indent=2) - added = len(truly_new) - print(f" ✅ 蒸馏新增 {added} 条") - - return all_new, added - - -# ── 主流程 ────────────────────────────────────────────── -def main(): - print("=" * 65) - print(" 小狸 AI - 最大化训练 v2.0") - print(f" 设备: {config.device} | AMP: {config.use_amp}") - print(f" 模型: {config.d_model}d, {config.n_layers}L, {config.n_heads}H, {config.d_ff}ff") - print("=" * 65) - - # 1. 加载词表 - vocab_path = os.path.join(config.save_dir, 'vocab.json') - if not os.path.exists(vocab_path): - print("❌ 未找到 vocab.json，请先运行 build_vocab.py") - return - with open(vocab_path, 'r', encoding='utf-8') as f: - vocab = json.load(f) - tokenizer_instance.char_to_id = vocab.copy() - tokenizer_instance.id_to_char = {v: k for k, v in vocab.items()} - print(f"词表: {len(vocab)} tokens") - - # 2. 创建模型 - model = LLMModel( - vocab_size=len(vocab), - d_model=config.d_model, - n_heads=config.n_heads, - n_layers=config.n_layers, - d_ff=config.d_ff, - max_len=config.max_length, - dropout=config.dropout, - use_memory=config.use_memory, - use_thinking=config.use_thinking, - use_moe=config.use_moe, - num_experts=config.num_experts, - top_k=config.top_k, - ) - params = sum(p.numel() for p in model.parameters()) - print(f"参数量: {params:,} ({params/1e6:.2f}M)") - - # 3. 预训练阶段（如果启用） - if config.pt_enable: - print("\n" + "=" * 65) - print(" 阶段一：通用中文预训练") - print("=" * 65) - - pt_dir = config.pt_data_dir - if not pt_dir or not os.path.exists(pt_dir): - pt_dir = os.path.join(config.save_dir, 'pretrain_data') - os.makedirs(pt_dir, exist_ok=True) - if not list(Path(pt_dir).glob("*")): - print(" 📥 下载预训练数据 (SkyPile 小样本)...") - _download_pretrain_data(pt_dir) - - try: - pt_dataset = PretrainDataset(pt_dir, vocab, config.pt_max_length) - pt_loader = DataLoader( - pt_dataset, - batch_size=config.pt_batch_size, - shuffle=True, - num_workers=config.num_workers, - pin_memory=torch.cuda.is_available(), - ) - except (FileNotFoundError, ValueError) as e: - print(f" ⚠️ 预训练数据不可用: {e}") - print(f" ⚠️ 跳过预训练阶段。请准备数据后放入: {pt_dir}") - pt_dataset = None - pt_loader = None - - if pt_loader is not None: - total_pt_steps = len(pt_loader) * config.pt_epochs - trainer_pt = MaxTrainer(model, config, stage="pretrain") - trainer_pt.build_scheduler(total_pt_steps) - - for epoch in range(1, config.pt_epochs + 1): - loss = trainer_pt.train_epoch( - pt_loader, epoch, config.pt_epochs, - stage_name="pretrain", log_every=config.log_every - ) - val_loss = trainer_pt.evaluate(pt_loader) - print(f" → 预训练 Epoch {epoch}/{config.pt_epochs} | " - f"Loss: {loss:.4f} | Val: {val_loss:.4f}") - - trainer_pt.save_model( - os.path.join(config.save_dir, f'pretrain_epoch_{epoch}.pth'), - vocab, epoch - ) - - print(f" ✅ 预训练完成") - - # 4. 阶段二：角色对齐 - print("\n" + "=" * 65) - print(" 阶段二：角色对齐 (QA 角色化训练)") - print("=" * 65) - - merged_path = os.path.join(config.save_dir, 'xiaoli_merged_data_v2.json') - if not os.path.exists(merged_path): - print(f"❌ 未找到合并数据: {merged_path}") - return - - with open(merged_path, 'r', encoding='utf-8') as f: - raw_data = json.load(f) - - pairs = [] - for item in raw_data: - if isinstance(item, dict): - q = item.get("question", "") or item.get("input", "") or item.get("prompt", "") - a = item.get("answer", "") or item.get("output", "") or item.get("response", "") - if q and a: - pairs.append((q, a)) - - random.seed(config.seed) - random.shuffle(pairs) - - split = int(len(pairs) * 0.9) - train_pairs = pairs[:split] - val_pairs = pairs[split:] - print(f"QA 数据: {len(pairs)} 条 (训练 {len(train_pairs)} + 验证 {len(val_pairs)})") - - train_ds = ConversationDataset(train_pairs, vocab, config.max_length) - train_loader = DataLoader( - train_ds, batch_size=config.align_batch_size, shuffle=True, - num_workers=config.num_workers, pin_memory=torch.cuda.is_available(), - ) - val_ds = ConversationDataset(val_pairs, vocab, config.max_length) - val_loader = DataLoader( - val_ds, batch_size=config.align_batch_size, shuffle=False, - num_workers=0, - ) - - distill_questions = [] - for qpath in ['exam_distill_questions.json', 'distill_questions_v3.json', 'distill_questions_v2.json']: - fp = os.path.join(config.save_dir, qpath) - if os.path.exists(fp): - with open(fp, 'r', encoding='utf-8') as f: - distill_questions = json.load(f) - print(f"蒸馏问题: {len(distill_questions)} 道 ({qpath})") - break - - total_align_steps = len(train_loader) * config.align_epochs - trainer_align = MaxTrainer(model, config, stage="align") - trainer_align.build_scheduler(total_align_steps) - - best_val_loss = float('inf') - total_distill_added = 0 - - for epoch in range(1, config.align_epochs + 1): - train_loss = trainer_align.train_epoch( - train_loader, epoch, config.align_epochs, - stage_name="align", log_every=config.log_every - ) - - val_loss = trainer_align.evaluate(val_loader) - marker = " ★" if val_loss < best_val_loss else "" - if val_loss < best_val_loss: - best_val_loss = val_loss - trainer_align.save_model( - os.path.join(config.save_dir, 'xiaoli_v2_best.pth'), vocab, epoch - ) - - print(f" → QA Epoch {epoch}/{config.align_epochs} | " - f"Loss: {train_loss:.4f} | Val: {val_loss:.4f}{marker} | " - f"LR: {trainer_align.optimizer.param_groups[0]['lr']:.2e}") - - if epoch % config.save_interval == 0: - trainer_align.save_model( - os.path.join(config.save_dir, f'llm_model_v2_resume_epoch_{epoch}.pth'), - vocab, epoch - ) - - if epoch % config.archive_interval == 0: - trainer_align.save_model( - os.path.join(config.save_dir, f'llm_model_v2_epoch_{epoch}.pth'), - vocab, epoch - ) - resumes = sorted(glob.glob(os.path.join(config.save_dir, 'llm_model_v2_resume_*.pth'))) - for old in resumes[:-config.keep_resume]: - try: - os.remove(old) - except: - pass - - if config.distill_enable and distill_questions and epoch % config.distill_every == 0: - _, added = run_distillation( - model, vocab, config.device, config.max_length, - distill_questions, config.save_dir - ) - total_distill_added += added - - trainer_align.save_model( - os.path.join(config.save_dir, 'xiaoli_v2_final.pth'), vocab, config.align_epochs - ) - - for f in glob.glob(os.path.join(config.save_dir, 'llm_model_v2_resume_*.pth')): - try: - os.remove(f) - except: - pass - - print(f"\n{'='*65}") - print(f" 训练完成！蒸馏新增: {total_distill_added} 条") - print(f" 最终模型: {os.path.join(config.save_dir, 'xiaoli_v2_final.pth')}") - print(f"{'='*65}") - - print("\n📝 快速测试:") - chat = LLMChat(model, vocab, config.device, config.max_length) - for q in ["你好", "你是谁", "1+1等于几", "讲个笑话"]: - ans = chat.generate_response(q, max_length=30, temperature=0.5) - print(f" Q: {q}") - print(f" A: {ans}") - print() - - -def _download_pretrain_data(target_dir: str): - import urllib.request - samples = [ - "天空是蓝色的，白云在天空中飘来飘去。小鸟在树上唱歌，花儿在阳光下开放。", - "小明是一个小学生，他每天都要去学校上课。今天老师教大家认识汉字，小明学会了'天'、'地'、'人'三个字。", - "春天来了，大地变绿了。小草从土里钻出来，树木长出了新叶子，到处都是生机勃勃的景象。", - "猫是一种常见的宠物，它们有着柔软的毛发和灵活的身体。猫咪喜欢晒太阳，也喜欢玩毛线球。", - "数学是一门重要的学科，我们生活中处处都用得到数学。比如买东西需要算钱，做菜需要称量材料。", - "学习编程可以培养逻辑思维能力，通过写代码，我们可以让计算机帮助我们完成各种任务。", - "中国有着悠久的历史和灿烂的文化，四大发明——造纸术、印刷术、火药、指南针，对世界文明做出了重要贡献。", - "保护环境人人有责，我们应该节约用水，减少使用塑料袋，做好垃圾分类，让地球变得更美好。", - "星星在夜空中闪烁，月亮有时圆有时缺。宇宙中有无数的星系，地球只是其中一颗小小的行星。", - "友谊是人生中宝贵的财富，好朋友会在你开心时和你一起笑，在你难过时给你安慰和鼓励。", - "人工智能是当今最前沿的科技领域之一，它让机器能够像人一样思考和学习，改变着我们的生活方式。", - "读书是获取知识的重要途径，通过阅读，我们可以了解世界各地的风土人情，学习前人的智慧和经验。", - "运动对身体健康非常重要，跑步、游泳、打篮球都是很好的运动方式，每天坚持运动可以让我们更有活力。", - "音乐是一种美妙的艺术形式，不同的乐器能演奏出不同的声音，组合在一起就形成了动听的旋律。", - "做饭是一门生活技能，学会做饭不仅能照顾好自己的饮食，还能给家人朋友带来美味的享受。", - "植物通过光合作用将阳光转化为能量，它们吸收二氧化碳，释放出氧气，为地球上的生命提供了重要的保障。", - "狗是人类最忠诚的朋友，它们聪明、勇敢、重情义，能够帮助人类看家护院、导盲、搜救等。", - "海洋覆盖了地球表面的百分之七十，海洋中有各种各样的生物，从微小的浮游生物到巨大的鲸鱼。", - "互联网让世界变成了一个地球村，人们可以通过网络随时随地获取信息、交流沟通、学习工作。", - "时间管理是一项重要的能力，合理安排时间可以提高工作效率，让我们有更多的时间做自己喜欢的事情。", - ] - out_path = os.path.join(target_dir, "pretrain_samples.txt") - with open(out_path, 'w', encoding='utf-8') as f: - for s in samples: - f.write(s + "\n") - - wiki_samples_path = os.path.join(target_dir, "wiki_samples.json") - wiki_data = [ - {"text": "孔子（公元前551年—公元前479年），名丘，字仲尼，是中国古代伟大的思想家、教育家，儒家学派的创始人。"}, - {"text": "太阳是太阳系的中心天体，直径为139.2万公里，表面温度约为5500摄氏度，地球上的生命离不开太阳的光和热。"}, - {"text": "长城是中国古代的军事防御工程，始建于春秋战国时期，总长度超过2.1万公里，被誉为世界七大奇迹之一。"}, - {"text": "DNA是脱氧核糖核酸的缩写，是生物遗传信息的载体，它决定了生物的性状和特征，呈双螺旋结构。"}, - {"text": "Python是一种广泛使用的高级编程语言，由吉多·范罗苏姆于1991年首次发布，以其简洁易读的语法而闻名。"}, - {"text": "水的化学式是H2O，在标准大气压下，水的沸点是100摄氏度，冰点是0摄氏度，水是生命之源。"}, - {"text": "贝多芬是德国著名的作曲家和钢琴家，他创作了《命运交响曲》《月光奏鸣曲》等不朽的音乐作品。"}, - {"text": "光合作用是植物、藻类和某些细菌利用光能将二氧化碳和水转化为有机物和氧气的过程。"}, - {"text": "人类基因组计划于1990年启动，2003年完成，成功绘制了人类基因组的完整序列图，对医学研究具有重要意义。"}, - {"text": "丝绸之路是古代连接中国与中亚、欧洲的贸易路线，不仅促进了商品交流，也推动了东西方文化的传播与融合。"}, - ] - with open(wiki_samples_path, 'w', encoding='utf-8') as f: - json.dump(wiki_data, f, ensure_ascii=False, indent=2) - - print(f" ✅ 已生成示例预训练数据: {len(samples) + len(wiki_data)} 条") - print(f" ⚠️ 这仅供流程测试！实际训练请替换为大规模中文语料") - print(f" 💡 建议: https://huggingface.co/datasets/SkyWork/SkyPile-150B") - - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/web_server.py b/web_server.py deleted file mode 100644 index b5b8976..0000000 --- a/web_server.py +++ /dev/null @@ -1,128 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -"""小狸v1.2 Web聊天服务器 - 当前训练中的模型""" -import sys, os, json, time -sys.path.insert(0, '/workspace/xiaoli-llm-chat') -import torch -from flask import Flask, request, jsonify, send_from_directory -from llm_chat_improved import LLMModel, LLMChat, build_vocabulary - -print("=" * 50) -print(" 小狸 v1.2 Web 服务加载中...") -print("=" * 50) - -device = 'cpu' -model_dir = '/workspace/xiaoli-llm-chat' - -# 构建当前vocab -with open(f'{model_dir}/xiaoli_merged_data_v2.json', 'r', encoding='utf-8') as f: - all_data = json.load(f) -pairs = [] -for d in all_data: - q = d.get('question','') or d.get('input','') or d.get('prompt','') or '' - a = d.get('answer','') or d.get('output','') or d.get('response','') or '' - if q and a: - pairs.append((q, a)) -vocab = build_vocabulary(pairs) -print(f"词汇表: {len(vocab)} | 数据: {len(all_data)}条") - -# 加载模型 (epoch 50 = 干净epoch 41 + vocab扩展) -ckpt_path = f'{model_dir}/llm_model_v2_resume_epoch_50.pth' -print(f"加载: {ckpt_path}") -ckpt = torch.load(ckpt_path, map_location=device) -old_vocab_size = ckpt.get('model_config', {}).get('vocab_size', 0) -cfg = ckpt['model_config'] -print(f"断点vocab: {old_vocab_size}") - -model = LLMModel( - vocab_size=len(vocab), - d_model=cfg['d_model'], n_heads=cfg['n_heads'], - n_layers=cfg['n_layers'], d_ff=cfg['d_ff'], - max_len=cfg['max_len'], dropout=cfg.get('dropout', 0.1), - use_memory=False, use_thinking=False, use_moe=False, - num_experts=1, top_k=1, -) - -# 迁移权重 -old_sd = ckpt['model_state_dict'] -new_sd = model.state_dict() -if old_vocab_size != len(vocab): - for key in new_sd: - if key in old_sd: - w = old_sd[key] - if w.shape == new_sd[key].shape: - new_sd[key] = w.clone() - elif len(w.shape) == 2 and w.shape[0] == old_vocab_size: - new_sd[key][:old_vocab_size] = w - elif 'freq' in key: - pass - else: - try: new_sd[key] = w.clone() - except: pass - model.load_state_dict(new_sd) -else: - model.load_state_dict(old_sd) - -model.to(device) -model.eval() -chat = LLMChat(model, vocab, device, max_length=cfg['max_len']) -print(f"✅ 模型加载完成！ | voc={len(vocab)} | max_len={cfg['max_len']}") - -app = Flask(__name__) - -@app.route('/') -def index(): - return send_from_directory(model_dir, 'chat.html') - -@app.route('/api/chat', methods=['POST']) -def api_chat(): - data = request.get_json() - question = data.get('question', '').strip() - temp = float(data.get('temperature', 0.2)) - max_len = int(data.get('max_length', 100)) - history = data.get('history', None) - - if not question: - return jsonify({'error': '问题不能为空'}) - - # 如果有history，转换为对话格式 - history_pairs = None - if history and isinstance(history, list) and len(history) > 0: - history_pairs = [(h['q'], h['a']) for h in history] - - try: - t0 = time.time() - answer = chat.generate_response(question, max_length=max_len, temperature=temp, history=history_pairs) - elapsed = time.time() - t0 - return jsonify({ - 'question': question, - 'answer': answer, - 'time': f'{elapsed:.1f}s', - 'temperature': temp, - }) - except Exception as e: - return jsonify({'error': str(e)}) - -@app.route('/api/info') -def model_info(): - return jsonify({ - 'name': '小狸 v1.2', - 'version': 'epoch 52 (架构扩展: 7.9M→21.3M, max_len=256)', - 'vocab_size': len(vocab), - 'max_length': cfg['max_len'], - 'd_model': cfg['d_model'], - 'layers': cfg['n_layers'], - 'heads': cfg['n_heads'], - 'd_ff': cfg['d_ff'], - 'data_size': f'{len(all_data):,}条', - 'training': '新架构训练中 epoch 1→200', - 'features': ['多轮对话(1,694条)', '长文本(758条)', '我不知道(1,140条)'], - }) - -if __name__ == '__main__': - print("\n" + "=" * 50) - print(" 小狸 v1.2 聊天服务已启动！") - print(" 访问地址: http://0.0.0.0:5000") - print(" 建议温度: 0.1 ~ 0.2") - print("=" * 50) - app.run(host='0.0.0.0', port=5000, debug=False, threaded=True) \ No newline at end of file diff --git a/xiaoli_next.py b/xiaoli_next.py deleted file mode 100644 index 209ec29..0000000 --- a/xiaoli_next.py +++ /dev/null @@ -1,422 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -XiaLi-Next 统一架构 v1.0 -========================= -同一套架构，从 20M 到 7B+ 通用 -核心设计：RoPE + SwiGLU + Pre-RMSNorm + GQA + Weight Tying - -用法: - from xiaoli_next import XiaLiConfig, XiaLiModel - - # 20M 超小模型 - cfg = XiaLiConfig.preset('20m') - model = XiaLiModel(cfg) - - # 1B 中等模型 - cfg = XiaLiConfig.preset('1b') - model = XiaLiModel(cfg) - - # 7B 大模型 - cfg = XiaLiConfig.preset('7b') - model = XiaLiModel(cfg) -""" - -import math -import json -from typing import Optional, Tuple - -import torch -import torch.nn as nn -import torch.nn.functional as F - - -# ════════════════════════════════════════════════════════════ -# 1. 配置系统 -# ════════════════════════════════════════════════════════════ - -class XiaLiConfig: - """统一配置 - 通过 preset 快速切换规模""" - - def __init__(self): - # 架构 - self.d_model: int = 320 - self.n_layers: int = 8 - self.n_heads: int = 8 - self.n_kv_heads: int = 4 # GQA: KV头数（=n_heads 就是MHA，=1 就是MQA） - self.d_ff: int = 0 # 0=自动 = int(d_model * 8/3 * 2) SwiGLU需要额外gate - self.max_seq_len: int = 2048 # 最大序列长度 - self.dropout: float = 0.0 # 大模型dropout=0 - self.rope_theta: float = 10000.0 # RoPE base frequency - self.tie_weights: bool = True # Weight Tying - self.vocab_size: int = 20007 - - # 训练 - self.norm_eps: float = 1e-5 - - @classmethod - def preset(cls, name: str) -> 'XiaLiConfig': - """按规模取配置""" - cfgs = { - # ── 超小模型 ── - '20m': dict(d_model=288, n_layers=6, n_heads=6, n_kv_heads=3, max_seq_len=1024), - '26m': dict(d_model=320, n_layers=8, n_heads=8, n_kv_heads=4, max_seq_len=2048), - # ── 百M级 ── - '100m': dict(d_model=512, n_layers=12, n_heads=8, n_kv_heads=4, max_seq_len=2048), - '300m': dict(d_model=768, n_layers=16, n_heads=12, n_kv_heads=4, max_seq_len=4096), - # ── 十亿级 ── - '500m': dict(d_model=1024, n_layers=24, n_heads=16, n_kv_heads=8, max_seq_len=8192), - '1b': dict(d_model=1536, n_layers=24, n_heads=16, n_kv_heads=8, max_seq_len=8192), - '3b': dict(d_model=2560, n_layers=32, n_heads=24, n_kv_heads=8, max_seq_len=8192), - # ── 十亿级+ ── - '7b': dict(d_model=4096, n_layers=32, n_heads=32, n_kv_heads=8, max_seq_len=8192, rope_theta=500000.0), - '13b': dict(d_model=5120, n_layers=40, n_heads=40, n_kv_heads=10, max_seq_len=8192, rope_theta=500000.0), - '70b': dict(d_model=8192, n_layers=80, n_heads=64, n_kv_heads=8, max_seq_len=8192, rope_theta=500000.0), - } - cfg = cls() - if name in cfgs: - for k, v in cfgs[name].items(): - setattr(cfg, k, v) - if cfg.d_ff == 0: - # SwiGLU 标准: 8/3 * d_model（参数量等价于 4x 标准 FFN） - cfg.d_ff = int(cfg.d_model * 8 / 3) - return cfg - - -# ════════════════════════════════════════════════════════════ -# 2. RMSNorm -# ════════════════════════════════════════════════════════════ - -class XiaLiRMSNorm(nn.Module): - """RMS Layer Normalization""" - def __init__(self, dim: int, eps: float = 1e-5): - super().__init__() - self.eps = eps - self.weight = nn.Parameter(torch.ones(dim)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - rms = torch.sqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) - return x / rms * self.weight - - -# ════════════════════════════════════════════════════════════ -# 3. RoPE（旋转位置编码） -# ════════════════════════════════════════════════════════════ - -def precompute_rope_freqs(dim: int, max_len: int, theta: float = 10000.0) -> Tuple[torch.Tensor, torch.Tensor]: - """预计算 RoPE 频率""" - freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) - t = torch.arange(max_len, dtype=torch.float32) - angles = torch.outer(t, freqs) # [max_len, dim//2] - cos = angles.cos().unsqueeze(0).unsqueeze(2) # [1, max_len, 1, dim//2] - sin = angles.sin().unsqueeze(0).unsqueeze(2) - return cos, sin - -def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: - """对 x 应用 RoPE。x: [batch, seq, n_heads, head_dim]""" - head_dim = x.size(-1) - # cos/sin 是每对维度一个值 (head_dim//2)，需要扩张到 head_dim - cos = torch.repeat_interleave(cos, 2, dim=-1) # [..., d//2] -> [..., d] - sin = torch.repeat_interleave(sin, 2, dim=-1) - x1 = x[..., :head_dim // 2] # [..., d//2] - x2 = x[..., head_dim // 2:] # [..., d//2] - rotated = torch.cat([-x2, x1], dim=-1) # [..., d] - return x * cos + rotated * sin - - -# ════════════════════════════════════════════════════════════ -# 4. GQA 注意力（分组查询注意力） -# ════════════════════════════════════════════════════════════ - -class GroupedQueryAttention(nn.Module): - """分组查询注意力 - 小模型=MHA，大模型节省KV Cache""" - - def __init__(self, config: XiaLiConfig): - super().__init__() - self.n_heads = config.n_heads - self.n_kv_heads = config.n_kv_heads # GQA: KV头数 - self.head_dim = config.d_model // config.n_heads - self.n_rep = self.n_heads // self.n_kv_heads # 每个KV头被几个Q头共享 - - # 投影（无bias） - self.wq = nn.Linear(config.d_model, config.n_heads * self.head_dim, bias=False) - self.wk = nn.Linear(config.d_model, self.n_kv_heads * self.head_dim, bias=False) - self.wv = nn.Linear(config.d_model, self.n_kv_heads * self.head_dim, bias=False) - self.wo = nn.Linear(config.n_heads * self.head_dim, config.d_model, bias=False) - self.attn_dropout = nn.Dropout(config.dropout) - - def forward( - self, - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - batch, seq_len, _ = x.shape - - # QKV 投影 - q = self.wq(x).view(batch, seq_len, self.n_heads, self.head_dim) - k = self.wk(x).view(batch, seq_len, self.n_kv_heads, self.head_dim) - v = self.wv(x).view(batch, seq_len, self.n_kv_heads, self.head_dim) - - # RoPE - cos_ = cos[:, :seq_len] - sin_ = sin[:, :seq_len] - q = apply_rope(q, cos_, sin_) - k = apply_rope(k, cos_, sin_) - - # GQA: KV 头扩展到 Q 头数 - if self.n_rep > 1: - k = k.repeat_interleave(self.n_rep, dim=2) - v = v.repeat_interleave(self.n_rep, dim=2) - - # 转置为 [batch, n_heads, seq, head_dim] - q = q.transpose(1, 2) - k = k.transpose(1, 2) - v = v.transpose(1, 2) - - # 注意力计算 - scale = self.head_dim ** -0.5 - attn = torch.matmul(q, k.transpose(-2, -1)) * scale - - if mask is not None: - attn = attn + mask[:, :, :seq_len, :seq_len] - - attn = F.softmax(attn, dim=-1, dtype=torch.float32).to(q.dtype) - attn = self.attn_dropout(attn) - - out = torch.matmul(attn, v) # [batch, n_heads, seq, head_dim] - out = out.transpose(1, 2).contiguous().view(batch, seq_len, -1) - - return self.wo(out) - - -# ════════════════════════════════════════════════════════════ -# 5. SwiGLU 前馈网络 -# ════════════════════════════════════════════════════════════ - -class SwiGLU(nn.Module): - """SwiGLU 前馈网络 - 比 ReLU FFN 质量更高""" - - def __init__(self, config: XiaLiConfig): - super().__init__() - # d_ff 已按 8/3*d_model 计算。SwiGLU: w1和w3同维度，w2降回 - hidden = config.d_ff - self.w1 = nn.Linear(config.d_model, hidden, bias=False) - self.w2 = nn.Linear(hidden, config.d_model, bias=False) - self.w3 = nn.Linear(config.d_model, hidden, bias=False) - self.dropout = nn.Dropout(config.dropout) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - # SwiGLU(x) = (SiLU(x @ W1) * (x @ W3)) @ W2 - return self.w2(self.dropout(F.silu(self.w1(x)) * self.w3(x))) - - -# ════════════════════════════════════════════════════════════ -# 6. Transformer 块 -# ════════════════════════════════════════════════════════════ - -class XiaLiBlock(nn.Module): - """Transformer 块 - Pre-RMSNorm + GQA + SwiGLU""" - - def __init__(self, config: XiaLiConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - - self.attention = GroupedQueryAttention(config) - self.feed_forward = SwiGLU(config) - - self.attn_norm = XiaLiRMSNorm(config.d_model, config.norm_eps) - self.ff_norm = XiaLiRMSNorm(config.d_model, config.norm_eps) - - def forward( - self, - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - h = x + self.attention(self.attn_norm(x), cos, sin, mask) - out = h + self.feed_forward(self.ff_norm(h)) - return out - - -# ════════════════════════════════════════════════════════════ -# 7. 主模型 -# ════════════════════════════════════════════════════════════ - -class XiaLiModel(nn.Module): - """ - XiaLi-Next 统一架构 - - 核心设计： - - RoPE 位置编码（无学习的 position embedding） - - Pre-RMSNorm（训练稳定） - - SwiGLU FFN（同等参数量更高质量） - - GQA 注意力（大模型省 KV Cache） - - Weight Tying（embedding 与输出层共享权重） - - 所有 Linear 无 bias（更干净） - - 从 20M 到 7B+ 统一架构 - """ - - def __init__(self, config: XiaLiConfig): - super().__init__() - self.config = config - - # 词嵌入 - self.embedding = nn.Embedding(config.vocab_size, config.d_model) - - # 预计算 RoPE - cos, sin = precompute_rope_freqs( - config.d_model // config.n_heads, - config.max_seq_len * 2, - config.rope_theta, - ) - self.register_buffer('rope_cos', cos, persistent=False) - self.register_buffer('rope_sin', sin, persistent=False) - - # Transformer 块 - self.blocks = nn.ModuleList([ - XiaLiBlock(config, i) for i in range(config.n_layers) - ]) - - # 最终归一化 - self.norm = XiaLiRMSNorm(config.d_model, config.norm_eps) - - # 输出层（与 embedding 共享权重） - self.fc_out = nn.Linear(config.d_model, config.vocab_size, bias=False) - - if config.tie_weights: - self.fc_out.weight = self.embedding.weight # Weight Tying - - self._init_weights() - - def _init_weights(self): - """初始化权重""" - for name, param in self.named_parameters(): - if param.dim() >= 2: - nn.init.normal_(param, mean=0.0, std=0.02) - else: - nn.init.zeros_(param) - - def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: - batch, seq_len = x.shape - - # 词嵌入 - h = self.embedding(x) - - # RoPE - cos = self.rope_cos[:, :seq_len] - sin = self.rope_sin[:, :seq_len] - - # Transformer 层 - for block in self.blocks: - h = block(h, cos, sin, mask) - - # 最终归一化 - h = self.norm(h) - - # 输出 - return self.fc_out(h) - - def generate_mask(self, seq_len: int) -> torch.Tensor: - """生成因果掩码 [1, 1, seq, seq]""" - mask = torch.full((seq_len, seq_len), float('-inf')) - mask = torch.triu(mask, diagonal=1) - return mask.unsqueeze(0).unsqueeze(0) - - @torch.no_grad() - def generate( - self, - input_ids: torch.Tensor, - max_new_tokens: int = 100, - temperature: float = 0.7, - top_k: int = 40, - eos_id: int = 3, - ) -> torch.Tensor: - """自回归生成""" - self.eval() - for _ in range(max_new_tokens): - # 截断到 max_seq_len - seq = input_ids if input_ids.size(1) <= self.config.max_seq_len else input_ids[:, -self.config.max_seq_len:] - - mask = self.generate_mask(seq.size(1)).to(seq.device) - logits = self.forward(seq, mask) - logits = logits[:, -1, :] / temperature - - # Top-k 采样 - if top_k > 0: - v, _ = torch.topk(logits, min(top_k, logits.size(-1))) - logits[logits < v[:, -1:]] = float('-inf') - - probs = F.softmax(logits, dim=-1) - next_token = torch.multinomial(probs, 1) - input_ids = torch.cat([input_ids, next_token], dim=-1) - - if next_token.item() == eos_id: - break - - return input_ids - - def count_params(self) -> dict: - """统计参数量""" - total = sum(p.numel() for p in self.parameters()) - trainable = sum(p.numel() for p in self.parameters() if p.requires_grad) - # 各模块 - emb = self.embedding.weight.numel() - attn = sum(p.numel() for n, p in self.named_parameters() if 'attention' in n) - ffn = sum(p.numel() for n, p in self.named_parameters() if 'feed_forward' in n) - norm = sum(p.numel() for n, p in self.named_parameters() if 'norm' in n) - return { - 'total': total, - 'trainable': trainable, - 'embedding': emb, - 'attention': attn, - 'ffn': ffn, - 'norm': norm, - 'total_m': total / 1e6, - } - - -# ════════════════════════════════════════════════════════════ -# 8. 与原有系统兼容的训练适配 -# ════════════════════════════════════════════════════════════ - -# 导出兼容名（方便从原系统切换） -XiaLiNextModel = XiaLiModel -XiaLiNextConfig = XiaLiConfig - - -# ════════════════════════════════════════════════════════════ -# 9. 快速测试 -# ════════════════════════════════════════════════════════════ - -if __name__ == "__main__": - print("=" * 65) - print(" XiaLi-Next 架构测试") - print("=" * 65) - - for name in ['20m', '26m', '100m']: - cfg = XiaLiConfig.preset(name) - try: - model = XiaLiModel(cfg) - stats = model.count_params() - print(f"\n [{name}] {cfg.d_model}d × {cfg.n_layers}L × {cfg.n_heads}H (KV={cfg.n_kv_heads})") - print(f" 参数量: {stats['total_m']:.2f}M") - print(f" 其中: emb={stats['embedding']/1e6:.2f}M, attn={stats['attention']/1e6:.2f}M, " - f"ffn={stats['ffn']/1e6:.2f}M") - - # 前向测试 - x = torch.randint(0, cfg.vocab_size - 1, (1, 64)) - try: - out = model(x) - print(f" 前向: ✅ {out.shape}") - except Exception as e: - print(f" 前向: ❌ {e}") - except Exception as e: - print(f"\n [{name}] ❌ {e}") - - print(f"\n{'='*65}") - print(f" ✅ XiaLi-Next 就绪！") - print(f" 从 20M 到 7B+ 统一架构") - print(f"{'='*65}") \ No newline at end of file