Spaces:
Running
Running
File size: 1,623 Bytes
3905c4a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
import torch
import torch.nn as nn
from torch.nn import functional as F
import math, time, os
from torch.utils.data import Dataset, DataLoader
import tiktoken
# from torch.cuda.amp import autocast, GradScaler
from torch.amp.autocast_mode import autocast
from torch.amp.grad_scaler import GradScaler
from tqdm import tqdm
from datasets import load_dataset
from components.model import GPTModel
from components.tokenizer import encode, decode, tokenizer
def decode(tokens):
return tokenizer.decode(tokens)
class TextDataset(Dataset):
def __init__(self, hf_dataset, block_size):
self.dataset = hf_dataset
# self.tokenizer = tokenizer
self.block_size = block_size
def __len__(self):
return len(self.dataset["train"])
def __getitem__(self, idx):
# Start with a random index sample
rand_idx = torch.randint(0, len(self.dataset["train"]), (1,)).item()
text = self.dataset["train"][rand_idx]["text"]
tokens = encode(text)
# Keep appending more samples if too short
while len(tokens) < self.block_size + 1:
next_idx = torch.randint(0, len(self.dataset["train"]), (1,)).item()
next_text = self.dataset["train"][next_idx]["text"]
tokens.extend(encode(" " + next_text))
# Prevent runaway growth
if len(tokens) > self.block_size * 2:
break
# Truncate to block_size + 1
tokens = torch.tensor(tokens[: self.block_size + 1])
x = tokens[: self.block_size]
y = tokens[1 : self.block_size + 1]
return x.long(), y.long()
|