import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS']
plt.rcParams['axes.unicode_minus'] = False
# 1. 准备一个简单的中文数据集
text = "深度学习是人工智能的一个重要分支，它让计算机能够从数据中学习。"
chars = list(set(text))  # 获取所有不重复的字符
print("语料库中的字符:", ''.join(chars))
# 创建字符到索引的映射
char_to_idx = {ch: i for i, ch in enumerate(chars)}
idx_to_char = {i: ch for i, ch in enumerate(chars)}
vocab_size = len(chars)
print(f"词汇表大小: {vocab_size}")
# 2. 创建训练数据 - 构建(输入, 目标)对
seq_length = 5  # 输入序列长度
inputs = []
targets = []
for i in range(len(text) - seq_length):
    input_seq = text[i:i + seq_length]
    target_char = text[i + seq_length]
    
    inputs.append([char_to_idx[ch] for ch in input_seq])
    targets.append(char_to_idx[target_char])
# 转换为PyTorch张量
X = torch.tensor(inputs, dtype=torch.long)
y = torch.tensor(targets, dtype=torch.long)
print(f"\n输入序列示例: '{text[:seq_length]}' -> 目标: '{text[seq_length]}'")
print(f"X的形状: {X.shape}")  # (样本数, 序列长度)
print(f"y的形状: {y.shape}")  # (样本数,)
# 3. 定义一个简单的模型
class SimpleCharModel(nn.Module):
    def __init__(self, vocab_size, embedding_dim=32, hidden_dim=64):
        super(SimpleCharModel, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.rnn = nn.RNN(embedding_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, vocab_size)
    
    def forward(self, x):
        # x形状: (batch_size, seq_length)
        embedded = self.embedding(x)  # (batch_size, seq_length, embedding_dim)
        _, hidden = self.rnn(embedded)  # hidden: (1, batch_size, hidden_dim)
        output = self.fc(hidden.squeeze(0))  # (batch_size, vocab_size)
        return output
# 初始化模型
model = SimpleCharModel(vocab_size)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)
print(f"\n模型结构:")
print(model)
# 4. 训练模型
print("\n开始训练...")
epochs = 200
for epoch in range(epochs):
    model.train()
    optimizer.zero_grad()
    
    # 前向传播
    outputs = model(X)
    loss = criterion(outputs, y)
    
    # 反向传播
    loss.backward()
    optimizer.step()
    
    if (epoch + 1) % 50 == 0:
        print(f'轮次 [{epoch+1}/{epochs}], 损失: {loss.item():.4f}')
# 5. 使用模型进行预测
def predict_next_chars(model, start_string, num_chars=10):
    model.eval()
    chars = [ch for ch in start_string]
    
    print(f"\n预测结果:")
    print(start_string, end="")
    
    for _ in range(num_chars):
        # 准备输入
        input_seq = [char_to_idx[ch] for ch in chars[-seq_length:]]
        input_tensor = torch.tensor([input_seq], dtype=torch.long)
        
        # 预测下一个字符
        with torch.no_grad():
            output = model(input_tensor)
            probabilities = torch.softmax(output, dim=1)
            predicted_idx = torch.argmax(probabilities, dim=1).item()
        
        # 获取预测的字符
        predicted_char = idx_to_char[predicted_idx]
        print(predicted_char, end="")
        chars.append(predicted_char)
    
    print()
# 测试预测
predict_next_chars(model, "深度学", 10)
predict_next_chars(model, "人工智", 10)
def visualize_prediction(model, input_string):
    model.eval()
    
    # 准备输入
    input_seq = [char_to_idx[ch] for ch in input_string]
    input_tensor = torch.tensor([input_seq], dtype=torch.long)
    
    # 获取预测概率
    with torch.no_grad():
        output = model(input_tensor)
        probabilities = torch.softmax(output, dim=1).squeeze().numpy()
    
    # 可视化
    chars = list(char_to_idx.keys())
    plt.figure(figsize=(12, 4))
    bars = plt.bar(range(len(chars)), probabilities)
    plt.xticks(range(len(chars)), chars, rotation=45)
    plt.title(f'输入 "{input_string}" 后下一个字符的概率分布')
    plt.ylabel('概率')
    
    # 标出概率最高的几个字符
    top_indices = probabilities.argsort()[-3:][::-1]
    for i, idx in enumerate(top_indices):
        plt.annotate(f'{probabilities[idx]:.3f}', 
                    xy=(idx, probabilities[idx]), 
                    xytext=(idx, probabilities[idx] + 0.02),
                    ha='center', fontsize=10,
                    arrowprops=dict(arrowstyle='->', color='red'))
        print(f'候选 {i+1}: "{chars[idx]}" - 概率: {probabilities[idx]:.3f}')
    
    plt.tight_layout()
    plt.show()
# 可视化一些预测
visualize_prediction(model, "深度学")
visualize_prediction(model, "人工智")
import math
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super(PositionalEncoding, self).__init__()
        
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        
        self.register_buffer('pe', pe)
    
    def forward(self, x):
        return x + self.pe[:x.size(0), :]
class SimpleTransformerModel(nn.Module):
    def __init__(self, vocab_size, d_model=64, nhead=4, num_layers=2):
        super(SimpleTransformerModel, self).__init__()
        self.d_model = d_model
        
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoder = PositionalEncoding(d_model)
        
        encoder_layers = nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward=128)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layers, num_layers)
        
        self.fc = nn.Linear(d_model, vocab_size)
    
    def forward(self, x):
        # x形状: (seq_length, batch_size)
        embedded = self.embedding(x) * math.sqrt(self.d_model)
        embedded = self.pos_encoder(embedded)
        
        output = self.transformer_encoder(embedded)
        output = self.fc(output[-1, :, :])  # 只取最后一个时间步
        return output
# 使用Transformer模型重新训练
print("\n=== 使用Transformer模型 ===")
transformer_model = SimpleTransformerModel(vocab_size)
optimizer = optim.Adam(transformer_model.parameters(), lr=0.005)
# 调整数据格式以适应Transformer
X_transformer = X.transpose(0, 1)  # Transformer需要(seq_len, batch_size)
# 简单训练几轮
for epoch in range(100):
    transformer_model.train()
    optimizer.zero_grad()
    
    outputs = transformer_model(X_transformer)
    loss = criterion(outputs, y)
    
    loss.backward()
    optimizer.step()
    
    if (epoch + 1) % 25 == 0:
        print(f'轮次 [{epoch+1}/100], 损失: {loss.item():.4f}')
# 测试Transformer模型
def transformer_predict(model, start_string, num_chars=10):
    model.eval()
    chars = [ch for ch in start_string]
    
    print(f"\nTransformer预测结果:")
    print(start_string, end="")
    
    for _ in range(num_chars):
        input_seq = [char_to_idx[ch] for ch in chars[-seq_length:]]
        input_tensor = torch.tensor([input_seq], dtype=torch.long).transpose(0, 1)
        
        with torch.no_grad():
            output = model(input_tensor)
            probabilities = torch.softmax(output, dim=1)
            predicted_idx = torch.argmax(probabilities, dim=1).item()
        
        predicted_char = idx_to_char[predicted_idx]
        print(predicted_char, end="")
        chars.append(predicted_char)
    
    print()
transformer_predict(transformer_model, "深度学", 10)