#循环神经网络 (RNN)：处理序列数据的逻辑

import torch
import torch.nn as nn

class RNNTextClassifier(nn.Module):
    """简单的情感二分类：输入是token id序列，输出是正/负的logits"""
    def __init__(
        self, 
        vocab_size: int,        # 词表大小
        embed_dim: int = 128,   # 词嵌入维度
        hidden_dim: int = 128,  # RNN隐藏层维度
        num_layers: int = 2,     # RNN堆叠层数
        num_classes: int = 2,    # 分类数
        dropout: float = 0.3     # 防止过拟合的dropout率
    ):
        super().__init__()
        # 1. 词嵌入层：把token id转换成稠密向量
        self.embedding = nn.Embedding(
            vocab_size, embed_dim, padding_idx=0  # padding_idx=0表示填充的向量不参与更新
        )
        # 2. 单向 RNN 层
        self.rnn = nn.RNN(
            input_size=embed_dim,
            hidden_size=hidden_dim,
            num_layers=num_layers,
            batch_first=True,  # 关键！输入输出的第一维是batch_size，方便处理
            dropout=dropout if num_layers > 1 else 0.0,  # 只有多层RNN时才加dropout
            bidirectional=False  # 单向，只看前文
        )
        # 3. 全连接分类层
        self.classifier = nn.Linear(hidden_dim, num_classes)

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        """
        input_ids: (batch_size, seq_len) → batch_size个句子，每个句子seq_len个token id
        return: (batch_size, num_classes) → 每个句子的正/负logits
        """
        # 步骤1：词嵌入
        embedded = self.embedding(input_ids)  # (batch_size, seq_len, embed_dim)
        
        # 步骤2：过RNN
        # output: (batch_size, seq_len, hidden_dim) → 每个时间步的隐藏状态
        # hidden: (num_layers, batch_size, hidden_dim) → 每个堆叠层最后时刻的隐藏状态
        output, hidden = self.rnn(embedded)
        
        # 步骤3：取RNN的最后一个时间步的隐藏状态做分类
        # batch_first=True时，-1就是序列的最后一个时刻
        last_hidden = output[:, -1, :]
        logits = self.classifier(last_hidden)
        return logits

# 测试一下
if __name__ == "__main__":
    vocab_size = 10000  # 假设词表有10000个词
    model = RNNTextClassifier(vocab_size=vocab_size)
    # 生成假数据：32个batch，每个句子50个token（id从1到9999，0是填充）
    dummy_input = torch.randint(1, vocab_size, (32, 50))
    dummy_logits = model(dummy_input)
    print(f"输出形状：{dummy_logits.shape}")  # 应该是 torch.Size([32, 2])

import torch
import torch.nn as nn

class BiRNNTextClassifier(nn.Module):
    """双向情感二分类：同时看前后文的语义"""
    def __init__(
        self, 
        vocab_size: int,
        embed_dim: int = 128,
        hidden_dim: int = 128,
        num_layers: int = 2,
        num_classes: int = 2,
        dropout: float = 0.3
    ):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        # 关键！把bidirectional设为True
        self.rnn = nn.RNN(
            input_size=embed_dim,
            hidden_size=hidden_dim,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout if num_layers > 1 else 0.0,
            bidirectional=True
        )
        # 因为是双向，隐藏维度要翻倍
        self.classifier = nn.Linear(hidden_dim * 2, num_classes)

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        embedded = self.embedding(input_ids)
        # RNN返回的hidden是(forward_hidden, backward_hidden)，每个的形状为(num_layers, batch_size, hidden_dim)
        output, (forward_hidden, backward_hidden) = self.rnn(embedded)
        # 取最后层的前向和后向隐藏状态，拼接
        combined_hidden = torch.cat([forward_hidden[-1], backward_hidden[-1]], dim=-1)
        logits = self.classifier(combined_hidden)
        return logits

# 同样测试一下
if __name__ == "__main__":
    vocab_size = 10000
    model = BiRNNTextClassifier(vocab_size=vocab_size)
    dummy_input = torch.randint(1, vocab_size, (32, 50))
    dummy_logits = model(dummy_input)
    print(f"双向RNN输出形状：{dummy_logits.shape}")  # 仍然是 torch.Size([32, 2])