1252 lines
48 KiB
Python
1252 lines
48 KiB
Python
#!/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()
|