Update pages/1_TensorIntro.py

pages/1_TensorIntro.py

@@ -207,27 +207,44 @@ print("Normalized data:", normalized_data)
         "code": '''import torch
 import torch.nn as nn
 import torch.optim as optim
-from torchtext.legacy import data, datasets
-
-# Define the fields for the dataset
-TEXT = data.Field(tokenize='spacy', include_lengths=True)
-LABEL = data.LabelField(dtype=torch.float)
-
-# Load the IMDb dataset
-train_data, test_data = datasets.IMDB.splits(TEXT, LABEL)
-
-# Build the vocabulary
-TEXT.build_vocab(train_data, max_size=25000)
-LABEL.build_vocab(train_data)
-
-# Create the iterators
-BATCH_SIZE = 64
-train_iterator, test_iterator = data.BucketIterator.splits(
-    (train_data, test_data),
-    batch_size=BATCH_SIZE,
-    sort_within_batch=True,
-    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-)
+from torchtext.datasets import IMDB
+from torchtext.data.utils import get_tokenizer
+from torchtext.vocab import build_vocab_from_iterator
+from torch.utils.data import DataLoader
+from torch.nn.utils.rnn import pad_sequence
+
+# Define the tokenizer and vocabulary
+tokenizer = get_tokenizer('basic_english')
+train_iter = IMDB(split='train')
+
+def yield_tokens(data_iter):
+    for _, text in data_iter:
+        yield tokenizer(text)
+
+vocab = build_vocab_from_iterator(yield_tokens(train_iter), specials=["<unk>"])
+vocab.set_default_index(vocab["<unk>"])
+
+# Define the text and label preprocessing pipeline
+text_pipeline = lambda x: vocab(tokenizer(x))
+label_pipeline = lambda x: 1 if x == 'pos' else 0
+
+# Define the collate function for the DataLoader
+def collate_batch(batch):
+    label_list, text_list, lengths = [], [], []
+    for _label, _text in batch:
+        label_list.append(label_pipeline(_label))
+        processed_text = torch.tensor(text_pipeline(_text), dtype=torch.int64)
+        text_list.append(processed_text)
+        lengths.append(processed_text.size(0))
+    label_list = torch.tensor(label_list, dtype=torch.float)
+    text_list = pad_sequence(text_list, batch_first=True)
+    lengths = torch.tensor(lengths, dtype=torch.int64)
+    return label_list, text_list, lengths
+
+# Create DataLoaders for training and testing
+train_iter, test_iter = IMDB()
+train_dataloader = DataLoader(list(train_iter), batch_size=8, shuffle=True, collate_fn=collate_batch)
+test_dataloader = DataLoader(list(test_iter), batch_size=8, shuffle=False, collate_fn=collate_batch)
 
 # Define the LSTM model
 class LSTM(nn.Module):
@@ -239,12 +256,12 @@ class LSTM(nn.Module):
 
     def forward(self, text, text_lengths):
         embedded = self.embedding(text)
-        packed_embedded = nn.utils.rnn.pack_padded_sequence(embedded, text_lengths)
+        packed_embedded = nn.utils.rnn.pack_padded_sequence(embedded, text_lengths, batch_first=True, enforce_sorted=False)
         packed_output, (hidden, cell) = self.lstm(packed_embedded)
         return self.fc(hidden.squeeze(0))
 
 # Instantiate the model
-INPUT_DIM = len(TEXT.vocab)
+INPUT_DIM = len(vocab)
 EMBEDDING_DIM = 100
 HIDDEN_DIM = 256
 OUTPUT_DIM = 1
@@ -258,11 +275,10 @@ optimizer = optim.Adam(model.parameters())
 N_EPOCHS = 5
 for epoch in range(N_EPOCHS):
     model.train()
-    for batch in train_iterator:
+    for labels, text, text_lengths in train_dataloader:
         optimizer.zero_grad()
-        text, text_lengths = batch.text
         predictions = model(text, text_lengths).squeeze(1)
-        loss = criterion(predictions, batch.label)
+        loss = criterion(predictions, labels)
         loss.backward()
         optimizer.step()