train_LM_from_scratch_DS_with_nlp.py

class NlpRawTextDataset(Dataset):
    def __init__(self, tokenizer, file_path: str, block_size: int):
        self.tokenizer = tokenizer
        self.file_path = file_path
        self.block_size = block_size
        print("Loading Dataset...")
        self.dataset = load_dataset("text", data_files=file_path)["train"]
        print("Loaded Dataset!")
        self.len = len(self.dataset)

    def __len__(self):
        return self.len

    def preprocess(self, text):
        batch_encoding = self.tokenizer(str(text), add_special_tokens=True, truncation=True, max_length=self.block_size)
        return torch.tensor(batch_encoding["input_ids"])

    def __getitem__(self, i):

        phrase = self.dataset[i]
        example = self.preprocess(phrase)
        return example