Implement a Tiny Feed-Forward Next-Token Predictor in PyTorch (Fill TODOs)

Task: Fill the TODOs to implement a tiny next-token prediction model in PyTorch (current-token conditioned) You are given a PyTorch skeleton and a synthetic dataset SyntheticDataset where each sample is a token sequence of length seq_len with tokens in [0, vocab_size-1]. Fill in all TODOs to implement and train a feed-forward only next-token prediction model. Constraints / Requirements The model must be a simple feed-forward network (no RNN/Transformer). Predicting token at t+1 must be conditioned only on the current token at t (i.e., per-position independent classification). You must complete: DataLoader TinyModel layers forward() input/target preparation for next-token prediction the training loop (forward / loss / backward / optimizer step) Use nn.CrossEntropyLoss(). Flatten all token positions in the batch to compute the loss. Skeleton Code import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, Dataset class SyntheticDataset(Dataset): def __init__(self, num_samples=1000, seq_len=10, vocab_size=100): torch.manual_seed(42) self.data = torch.randint(0, vocab_size, (num_samples, seq_len)) def __len__(self): return len(self.data) def __getitem__(self, idx): return self.data[idx] vocab_size = 100 dataset = SyntheticDataset(vocab_size=vocab_size) # TODO: Create dataloader with batch size of 32 data_loader = None class TinyModel(nn.Module): """Embedding -> Linear -> ReLU -> Linear""" def __init__(self, vocab_size, embed_dim=32, hidden_dim=64): super().__init__() # TODO: Define the layers. pass def forward(self, x): # TODO: Implement the forward pass pass def train_model(model, dataloader, epochs=3, learning_rate=0.001): criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=learning_rate) model.train() for epoch in range(epochs): total_epoch_loss = 0 for batch_idx, sequences in enumerate(dataloader): # TODO: # 1) Prepare inputs/targets for next-token prediction # 2) Flatten positions # 3) forward, compute loss # 4) backward and optimizer step pass avg_epoch_loss = total_epoch_loss / len(dataloader) print(f"Epoch {epoch+1}/{epochs}, Loss: {avg_epoch_loss:.4f}") my_model = TinyModel(vocab_size) train_model(my_model, data_loader) Expected Behavior (must match) data_loader = DataLoader(dataset, batch_size=32, shuffle=True) Model: Embedding(vocab_size, embed_dim) → Linear(embed_dim, hidden_dim) → ReLU → Linear(hidden_dim, vocab_size) For a batch sequences of shape [B, L]: inputs = sequences[:, :-1] → shape [B, L-1] targets = sequences[:, 1:] → shape [B, L-1] Flatten positions so that CrossEntropyLoss sees: logits: [B*(L-1), vocab_size] targets: [B*(L-1)] Training loop must include optimizer.zero_grad(), loss.backward(), optimizer.step(), and epoch loss accumulation. Example Input (no stdin; run script as-is) Output Should print 3 lines like: `Epoch 1/3, Loss: ...` and loss should be a finite float.