Токенизация больших наборов текстовых данных

Question

Я пытаюсь работать над проектом генерации текста. Я скачал набор данных WikiBooks с Kaggle:

https://www.kaggle.com/datasets/dhruvildave/wikibooks-dataset

И когда я пытаюсь создать набор данных для токенизации текстов, мое ядро ​​выходит из строя, потому что ему не хватает памяти.

      type from torch.utils.data import Dataset
from transformers import BertTokenizer

import os

class WikiDataset(Dataset):
  def __init__(self, query, conn, tokenizer, max_len):
    self.df = pd.read_sql_query(query, conn)
    self.tokenizer = tokenizer
    self.max_len = max_len
  
  def __len__(self):
    return len(self.df)

  def __getitem__(self, index):
    # Return tokenized sequence
    texts = self.df[index]

    tokens = self.tokenizer.encode_plus(
        texts,
        max_length=self.max_len,
        truncation=True,
        padding="max_length",
        return_tensors="pt"
    )

    return tokens["input_ids"], tokens["attention_masks"]

tokenizer = BertTokenizer.from_pretrained("bert-base-cased")
train_ds = WikiDataset(query, conn, tokenizer, 128)

Мой вопрос: как токенизировать тексты без сбоя ядра?

Я пробовал разные значения «max_length», но предполагаю, что основная проблема — размер набора данных.