3

我想在列格式的训练数据上训练一个 PyTorch NLP 模型，我想用pyspark 数据框Dataset作为原始数据来构建一个 PyTorch（不确定它是不是正确的方法......）。

为了预处理文本，我使用了transformers库提供的分词器和tokenizing_UDF应用分词的函数。

然后将该Dataset对象馈送到 aDataLoader以训练 ML 模型。

我目前拥有的是这样的：

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()以便我TokenizedDataset可以处理熊猫数据框。

这是一个明智的做法吗？如果是这样，我应该如何修改TokenizedDataset代码以直接处理 pyspark 数据帧？如果我偏离了轨道，我应该改用https://github.com/uber/petastorm吗？

4

0 回答 0