import torch
import torch.nn as nn
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

class SentimentLSTM(nn.Module):
    def __init__(self, vocab_size, embed_dim=128, hidden_dim=256,
                 num_layers=2, dropout=0.3, num_classes=2):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.lstm = nn.LSTM(
            embed_dim, hidden_dim,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout,
            bidirectional=True
        )
        self.dropout = nn.Dropout(dropout)
        # bidirectional -> hidden_dim * 2
        self.classifier = nn.Linear(hidden_dim * 2, num_classes)

    def forward(self, token_ids, lengths):
        # token_ids: [batch, max_seq_len]
        embedded = self.dropout(self.embedding(token_ids))

        # Pack sequences to skip PAD tokens in LSTM computation
        packed = pack_padded_sequence(embedded, lengths,
                                      batch_first=True, enforce_sorted=False)
        packed_out, (h_n, c_n) = self.lstm(packed)
        output, _ = pad_packed_sequence(packed_out, batch_first=True)

        # Use last non-padded hidden states (from both directions)
        # h_n shape: [num_layers * 2, batch, hidden_dim]
        # Take last layer: forward (index -2) and backward (index -1)
        forward_h = h_n[-2]   # [batch, hidden_dim]
        backward_h = h_n[-1]  # [batch, hidden_dim]
        combined = torch.cat([forward_h, backward_h], dim=1)

        return self.classifier(self.dropout(combined))

# Training
model = SentimentLSTM(vocab_size=25000)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
RNNs & LSTMs — רשתות רקורנטיביות וזיכרון לטווח ארוך

RNNs & LSTMs — רשתות רקורנטיביות וזיכרון לטווח ארוך
