Upload folder using huggingface_hub

utils/data_loader.py

@@ -0,0 +1,77 @@
+import torch
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+
+def get_mnist_loaders(batch_size=64):
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((0.1307,), (0.3081,))
+    ])
+
+    train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
+    test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
+
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
+
+    return train_loader, test_loader
+
+def get_fashion_mnist_loaders(batch_size=64):
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((0.5,), (0.5,))
+    ])
+
+    train_dataset = datasets.FashionMNIST(root='./data', train=True, download=True, transform=transform)
+    test_dataset = datasets.FashionMNIST(root='./data', train=False, download=True, transform=transform)
+
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
+
+    return train_loader, test_loader
+
+def get_imdb_loaders(batch_size=64, max_len=256, vocab_size=10000):
+    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, Dataset
+    import torch.nn.utils.rnn as rnn_utils
+
+    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>", "<pad>"])
+    vocab.set_default_index(vocab["<unk>"])
+
+    def text_pipeline(text):
+        return vocab(tokenizer(text))
+
+    class IMDBDataset(Dataset):
+        def __init__(self, split):
+            self.data = list(IMDB(split=split))
+            self.max_len = max_len
+
+        def __len__(self):
+            return len(self.data)
+
+        def __getitem__(self, idx):
+            label, text = self.data[idx]
+            # Convert label: 1 (neg), 2 (pos) -> 0, 1
+            label = 0 if label == 1 else 1
+            tokens = text_pipeline(text)[:self.max_len]
+            # Padding
+            if len(tokens) < self.max_len:
+                tokens += [vocab["<pad>"]] * (self.max_len - len(tokens))
+            return torch.tensor(tokens), torch.tensor(label)
+
+    train_dataset = IMDBDataset('train')
+    test_dataset = IMDBDataset('test')
+
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
+
+    return train_loader, test_loader, len(vocab)