Должен ли я создать набор данных PyTorch для обучения модели с использованием фрейма данных pyspark?

Question

Я хочу обучить модель PyTorch NLP над обучающими данными в столбчатом формате, и я подумал создать PyTorch, используя в качестве необработанных данных фрейм данных pyspark (не уверен, что это правильный подход ...).

Для предварительной обработки текста я использую токенизатор, предоставленный transformers библиотека и tokenizing_UDF функция для применения токенизации.

В Dataset затем объект передается в DataLoader обучить модель машинного обучения.

Что у меня сейчас есть:

      import pandas as pd # ideally I'd like to get rid of pandas here
import torch
from torch.utils.data.dataset import Dataset
from transformers import BertTokenizer
from pyspark.sql import types as T
from pyspark.sql import functions as F

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

text = ["This is a test.", "This is not a test."]*100
label = [1, 0]*100

df = sqlContext.createDataFrame(zip(text, label), schema=['text', 'label'])
tokenizing_UDF = udf(lambda t: tokenizer.encode(t),  T.ArrayType(T.LongType())) 
df = df.withColumn("tokenized", tokenizing_UDF(F.col("text"))) # not sure this is the right way
df = df.toPandas() # ugly

class TokenizedDataset(Dataset):
    """needs refactoring..."""
    def __init__(self, df):
        self.data = df
        
    def __getitem__(self, index):
        text = self.data.loc[index].tokenized
        text = torch.LongTensor(text)
        label = self.data.loc[index].label
        return (text, label)

    def __len__(self):
        count = len(self.data)
        return count

dataset = TokenizedDataset(df) # slow...

Я сейчас вызываю .toPandas() так что я могу иметь дело с фреймом данных pandas.

Это разумный подход? Если да, то как мне изменить TokenizedDatasetкод для обработки фреймов данных pyspark напрямую? Если я сбился с пути, следует ли мне использовать https://github.com/uber/petastorm вместо этого?