main.py

import json
import math
import os
import jieba
import numpy as np

# ==============================================================================
# 0. 全局可调整参数 (Hyperparameters & Configs)
# ==============================================================================
EMBED_DIM = 128  # 词向量维度
HIDDEN_DIM = 512  # 隐藏层维度
LEARNING_RATE = 0.05  # 学习率
EPOCHS = 1000  # 训练轮数
BATCH_SIZE = 100  # 新增：每批处理的句子数量
PRINT_EVERY = 100  # 每隔多少轮打印一次日志

VOCAB_PATH = 'vocab.json'
CORPUS_PATH = 'datasets.jsonl'
MODEL_PATH = 'generative_model.npz'  # 使用 .npz 替代 .json，高效二进制存储


# ==============================================================================
# 1. 环境准备与数据加载函数
# ==============================================================================
def prepare_mock_data():
    """生成模拟的 JSONL 语料和词表（仅用于演示）"""
    if not os.path.exists(CORPUS_PATH):
        mock_corpus = [
            {"text": "苹果的总部在加利福尼亚。"},
            {"text": "谷歌专注于搜索引擎和人工智能。"},
            {"text": "微软开发了微软Windows操作系统。"}
        ]
        with open(CORPUS_PATH, 'w', encoding='utf-8') as f:
            for item in mock_corpus:
                f.write(json.dumps(item, ensure_ascii=False) + '\n')
        print(f"[提示]: 已自动创建模拟语料文件: {CORPUS_PATH}")

    if not os.path.exists(VOCAB_PATH):
        vocab = {}
        idx = 0
        with open(CORPUS_PATH, 'r', encoding='utf-8') as f:
            for line in f:
                item = json.loads(line.strip())
                for word in jieba.cut(item['text']):
                    if word not in vocab:
                        vocab[word] = idx
                        idx += 1
        with open(VOCAB_PATH, 'w', encoding='utf-8') as f:
            json.dump(vocab, f, ensure_ascii=False)
        print(f"[提示]: 已自动创建模拟词表文件: {VOCAB_PATH}")


def load_vocab(path):
    with open(path, 'r', encoding='utf-8') as f:
        return json.load(f)


def read_jsonl_corpus(path):
    """读取 jsonl 文件中的语料，返回句子列表"""
    sentences = []
    with open(path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if line:
                data = json.loads(line)
                if 'text' in data:
                    sentences.append(data['text'])
    return sentences


# ==============================================================================
# 2. 模块化模型类定义
# ==============================================================================
class SimpleLanguageModel:
    def __init__(self, vocab_size, embed_dim, hidden_dim):
        self.vocab_size = vocab_size
        self.embed_dim = embed_dim
        self.hidden_dim = hidden_dim

        # 初始化权重
        self.word_vectors = np.random.randn(vocab_size, embed_dim) * 0.01
        self.W1 = np.random.randn(embed_dim, hidden_dim) * 0.1
        self.b1 = np.zeros(hidden_dim)
        self.W2 = np.random.randn(hidden_dim, vocab_size) * 0.1
        self.b2 = np.zeros(vocab_size)

    def forward(self, batch_word_ids):
        """
        前向传播（支持批量）
        batch_word_ids: 包含多个句子 ID 列表的列表，例如 [[1, 2], [3, 4, 5]]
        """
        batch_size = len(batch_word_ids)

        # 因为各句子长度不等，这里手动对每个句子做平均池化，拼接成 (batch_size, embed_dim)
        sentence_vectors = np.zeros((batch_size, self.embed_dim))
        for i, word_ids in enumerate(batch_word_ids):
            sentence_vectors[i] = np.mean(self.word_vectors[word_ids], axis=0)

        # 隐藏层 + ReLU -> 维度: (batch_size, hidden_dim)
        h_linear = np.dot(sentence_vectors, self.W1) + self.b1
        h = np.maximum(0, h_linear)

        # 输出层 -> 维度: (batch_size, vocab_size)
        logits = np.dot(h, self.W2) + self.b2

        # Softmax (按行/按样本计算)
        shift_logits = logits - np.max(logits, axis=1, keepdims=True)
        exp_logits = np.exp(shift_logits)
        probs = exp_logits / np.sum(exp_logits, axis=1, keepdims=True)

        return sentence_vectors, h, probs

    def backward(self, batch_word_ids, sentence_vectors, h, probs, target_ids, lr):
        """
        反向传播与参数更新（支持批量）
        target_ids: 当前批次每个样本的目标标签 ID 数组，形状为 (batch_size,)
        """
        batch_size = len(batch_word_ids)

        # 1. 输出层梯度 d_output -> 维度: (batch_size, vocab_size)
        d_output = probs.copy()
        d_output[np.arange(batch_size), target_ids] -= 1
        d_output /= batch_size  # 取批次平均

        # 2. W2 和 b2 梯度
        dW2 = np.dot(h.T, d_output)
        db2 = np.sum(d_output, axis=0)

        # 3. 隐藏层梯度
        d_hidden = np.dot(d_output, self.W2.T)
        d_hidden[h <= 0] = 0

        # 4. W1 和 b1 梯度
        dW1 = np.dot(sentence_vectors.T, d_hidden)
        db1 = np.sum(d_hidden, axis=0)

        # 5. 词向量梯度（由于是不规则的列表结构，针对每个样本单独回传）
        word_grad = np.zeros_like(self.word_vectors)
        for i, word_ids in enumerate(batch_word_ids):
            # 获取当前句子对 Embedding 的梯度
            d_sentence = np.dot(d_hidden[i], self.W1.T) / len(word_ids)
            for wid in word_ids:
                word_grad[wid] += d_sentence

        # 6. 梯度下降更新参数
        self.word_vectors -= lr * word_grad
        self.W1 -= lr * dW1
        self.b1 -= lr * db1
        self.W2 -= lr * dW2
        self.b2 -= lr * db2

    def train_on_corpus(self, sentences, vocab, epochs, batch_size, lr, print_every):
        """在给定的语料上进行批量训练"""
        print(f"\n开始 Batch 训练，共 {epochs} 轮，Batch Size = {batch_size}...")

        # 过滤并预转义所有句子为 ID 列表
        tokenized_corpus = []
        for s in sentences:
            ids = [vocab[w] for w in jieba.cut(s) if w in vocab]
            if ids:
                tokenized_corpus.append(ids)

        if not tokenized_corpus:
            print("错误: 语料库中没有词表内可识别的词，无法训练。")
            return

        num_sentences = len(tokenized_corpus)

        for epoch in range(1, epochs + 1):
            # 随机采样一个 batch 的句子索引
            batch_indices = np.random.choice(num_sentences, size=batch_size, replace=True)

            batch_word_ids = [tokenized_corpus[idx] for idx in batch_indices]
            # 随机选每个句子中的一个词作为预测目标
            target_ids = np.array([np.random.choice(word_ids) for word_ids in batch_word_ids])

            # 前向 + 反向
            sentence_vectors, h, probs = self.forward(batch_word_ids)

            # 计算 Batch 的平均损失
            loss = -np.mean(np.log(probs[np.arange(batch_size), target_ids] + 1e-15))

            self.backward(batch_word_ids, sentence_vectors, h, probs, target_ids, lr)

            if epoch % print_every == 0 or epoch == 1:
                print(f"第 {epoch:4d} 轮训练 -> 当前平均损失值(Loss): {loss:.6f}")

    def save_model(self, path):
        np.savez_compressed(
            path,
            word_vectors=self.word_vectors,
            W1=self.W1,
            b1=self.b1,
            W2=self.W2,
            b2=self.b2
        )
        print(f"[系统提示]: 模型已成功高效导出为二进制文件 '{path}'")

    def load_model(self, path):
        if os.path.exists(path):
            data = np.load(path)
            self.word_vectors = data['word_vectors']
            self.W1 = data['W1']
            self.b1 = data['b1']
            self.W2 = data['W2']
            self.b2 = data['b2']
            print(f"[系统提示]: 已成功从 '{path}' 加载模型参数。")
        else:
            print(f"[警告]: 未找到模型文件 '{path}'，保持当前初始化权重。")


# ==============================================================================
# 3. 主程序执行流程
# ==============================================================================
if __name__ == "__main__":
    prepare_mock_data()

    vocab = load_vocab(VOCAB_PATH)
    vocab_size = len(vocab)
    sentences = read_jsonl_corpus(CORPUS_PATH)

    print(f"成功加载词表，大小为: {vocab_size}")
    print(f"成功从 JSONL 加载了 {len(sentences)} 条语料。")

    model = SimpleLanguageModel(
        vocab_size=vocab_size,
        embed_dim=EMBED_DIM,
        hidden_dim=HIDDEN_DIM
    )

    # 训练模型（传入了 BATCH_SIZE）
    model.train_on_corpus(
        sentences=sentences,
        vocab=vocab,
        epochs=EPOCHS,
        batch_size=BATCH_SIZE,
        lr=LEARNING_RATE,
        print_every=PRINT_EVERY
    )

    model.save_model(MODEL_PATH)

    # --- 验证：尝试重新加载模型并做一次推理 ---
    print("\n" + "=" * 15 + " 重新加载并验证 " + "=" * 15)
    new_model = SimpleLanguageModel(vocab_size, EMBED_DIM, HIDDEN_DIM)
    new_model.load_model(MODEL_PATH)

    test_sentence = "苹果微软"
    test_ids = [vocab[w] for w in jieba.cut(test_sentence) if w in vocab]

    # 推理时，我们包裹成一个大小为 1 的 batch 传进去: [test_ids]
    _, _, final_probs = new_model.forward([test_ids])
    predicted_idx = np.argmax(final_probs[0])

    id_to_word = {v: k for k, v in vocab.items()}
    print(f"[验证结果]: 输入: '{test_sentence}'")
    print(f"             模型预测概率最高的下一个词是: '{id_to_word[predicted_idx]}'")