From f7e61f26d170cf25e23dec6af589c927c46f1d8a Mon Sep 17 00:00:00 2001
From: shuiyu1123 <f17376088816@163.com>
Date: Fri, 5 Jun 2026 15:08:42 +0000
Subject: [PATCH] =?UTF-8?q?=E4=B8=8A=E4=BC=A0=20llm=5Fchat=5Fimproved.py?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 llm_chat_improved.py | 1251 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 1251 insertions(+)
 create mode 100644 llm_chat_improved.py

diff --git a/llm_chat_improved.py b/llm_chat_improved.py
new file mode 100644
index 0000000..ace8095
--- /dev/null
+++ b/llm_chat_improved.py
@@ -0,0 +1,1251 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+LLM对话AI训练代码 - 改进版
+基于Transformer架构，参考MiniMind/nanoGPT等开源项目优化
+"""
+
+import torch
+import torch.nn as nn
+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 = {
+            '<PAD>': 0, '<UNK>': 1, '<SOS>': 2, '<EOS>': 3,
+            '<START>': 4, '<END>': 5, '<SEP>': 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)
+    
+    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
+        
+        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['<UNK>']]
+        
+        chars = self._split_chars(text)
+        ids = []
+        for char in chars:
+            if char in self.char_to_id:
+                ids.append(self.char_to_id[char])
+            else:
+                # 未知字符用 UNK token
+                ids.append(self.special_tokens['<UNK>'])
+        
+        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]
+                if char not in self.special_tokens:
+                    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)
+        
+        # 计算注意力分数
+        scores = torch.matmul(xq, xk.transpose(-2, -1)) / math.sqrt(self.d_head)
+        
+        # 应用掩码
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, float('-inf'))
+        
+        # Softmax
+        attn_weights = torch.softmax(scores, dim=-1)
+        attn_weights = self.dropout(attn_weights)
+        
+        # 注意力输出
+        output = torch.matmul(attn_weights, xv)
+        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 = 128):
+        self.conversations = conversations
+        self.vocab = vocab
+        self.max_length = max_length
+        self.pad_token_id = vocab.get('<PAD>', 0)
+        self.unk_token_id = vocab.get('<UNK>', 1)
+    
+    def __len__(self):
+        return len(self.conversations)
+    
+    def __getitem__(self, idx):
+        question, answer = self.conversations[idx]
+        
+        # 创建训练样本: <SOS> 问题 <SEP> <START> 答案 <END> <EOS>
+        text = f"<SOS>{question}<SEP><START>{answer}<END><EOS>"
+        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', use_amp: bool = None):
+        self.model = model.to(device)
+        self.device = device
+        self.criterion = nn.CrossEntropyLoss(ignore_index=0)
+        self.optimizer = optim.AdamW(model.parameters(), lr=0.0001, 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) // 10)  # 每10%打印一次
+        
+        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)
+            
+            # 前向传播
+            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 save_model(self, filepath: str, vocab: dict, epoch: int = None):
+        save_dict = {
+            'model_state_dict': self.model.state_dict(),
+            'optimizer_state_dict': self.optimizer.state_dict(),
+            'vocab': vocab,
+            'model_config': {
+                'vocab_size': self.model.vocab_size,
+                'd_model': self.model.d_model,
+                'n_heads': self.model.n_heads,
+                'n_layers': self.model.n_layers,
+                'd_ff': self.model.d_ff,
+                'max_len': self.model.max_len,
+                'use_memory': self.model.use_memory,
+                'use_thinking': self.model.use_thinking,
+                'use_moe': self.model.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)
+        self.model.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 = 128,
+                 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()}
+    
+    def generate_response(self, question: str, max_length: int = 100, temperature: float = 0.8,
+                         include_thinking: bool = False, use_beam_search: bool = False):
+        """生成回答"""
+        self.model.eval()
+        
+        # 使用与训练数据一致的格式: <SOS>问题<SEP><START>
+        input_text = f"<SOS>{question}<SEP><START>"
+        
+        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)
+                
+                # 获取最后一个时间步的预测
+                next_token_logits = outputs[0, -1, :] / 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('<END>', -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('<END>', -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('<SOS>', '').replace('<EOS>', '')
+        response = response.replace('<START>', '').replace('<END>', '')
+        response = response.replace('<SEP>', '')
+        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()