Update pages/17_RNN_News.py

pages/17_RNN_News.py

@@ -6,6 +6,7 @@ from torchtext.data.utils import get_tokenizer
 from torchtext.vocab import build_vocab_from_iterator
 from torchtext.datasets import AG_NEWS
 from torch.utils.data import DataLoader, random_split
+from torch.nn.utils.rnn import pad_sequence
 import matplotlib.pyplot as plt
 import pandas as pd
 import numpy as np
@@ -28,12 +29,11 @@ class RNN(nn.Module):
 
 # Create a custom collate function to pad sequences
 def collate_batch(batch):
-    label_list, text_list, lengths = [], [], []
+    label_list, text_list = [], []
     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))
     labels = torch.tensor(label_list, dtype=torch.int64)
     texts = pad_sequence(text_list, batch_first=True, padding_value=vocab["<pad>"])
     return texts, labels
@@ -191,4 +191,4 @@ if 'trained_model' in st.session_state and st.sidebar.button('Show Test Results'
     for i, (text, true_label, predicted) in enumerate(samples):
         st.write(f'Sample {i+1}')
         st.text(' '.join([vocab.get_itos()[token] for token in text]))
-        st.write(f'Ground Truth: {LABEL.vocab.itos[true_label.item()]}, Predicted: {LABEL.vocab.itos[predicted.item()]}')
+        st.write(f'Ground Truth: {true_label.item()}, Predicted: {predicted.item()}')