PySpark/Petastorm: как сгенерировать и включить элемент nn.Embedding()?

Question

Мне интересно, знает ли кто-нибудь, как получить доступ или сгенерировать новые данные при использовании Uber petastorm?

Что я сделал:

Используемый vectorAssembler на моих непрерывных функциях, чтобы создать искровой DF (преобразованный в petastorm) с 6 столбцами: характеристики, метки и 4 категориальные характеристики.

Затем я попытался сгенерировать и вернуть свои вложения через TransformSpec вот так:

def transform_row_train(pd_batch):
  """
  The input and output of this function must be pandas dataframes.
  Do data augmentation for the training dataset only.
  """
  # normalize features
  pd_batch['features'] = pd_batch['features'].map(lambda x: x - train_means.flatten())
  pd_batch['features'] = pd_batch['features'].map(lambda x: np.true_divide(x, train_stds.flatten() + 0.00001))
 
  from sklearn.preprocessing import LabelEncoder
  import torch
  import torch.nn as nn
  # grab categorical vars
  cat_df = pd_batch[['TAIL_NUM', 'OP_CARRIER_AIRLINE_ID', 'ORIGIN_AIRPORT_ID', 'DEST_AIRPORT_ID']]

  # convert to int64
  label_encoders = {}
  for cat_col in cat_df:
      label_encoders[cat_col] = LabelEncoder()
      cat_df[cat_col] = label_encoders[cat_col].fit_transform(cat_df[cat_col])
  # get number of unique cats
  cat_dims = [int(cat_df[col].nunique()) for col in cat_df]
  # generate embedding dims
  emb_dims = [(x, min(1, (x + 2) // 2)) for x in cat_dims]
  # generate embedding lookups
  all_embeddings = nn.ModuleList([nn.Embedding(ni, nf) for ni, nf in emb_dims])
  # grab the categorical data to interate over
  cats = cat_df.values
  cats = torch.tensor(cats).long()

  # create embeddings
  embeddings = []
  for i, emb in enumerate(all_embeddings):
      embeddings.append(emb(cats[:, i]))
  embeddings = torch.cat(embeddings, 1)
  
  return pd_batch, embeddings


def get_transform_spec_train():
  # Note that the output shape of the `TransformSpec` is not automatically known by petastorm, 
  # so we need to specify the shape for new columns in `edit_fields` and specify the order of 
  # the output columns in `selected_fields`.
  # 17 feature columns
  return TransformSpec(partial(transform_row_train),
                       #edit_fields=[('features', np.float32, (), False)],
                       #selected_fields=['DEP_DEL15', 'features', 'cats', 'cats1', 'cats2', 'cats3', 'cats4', 'cats5'])
                      removed_fields=['TAIL_NUM', 'OP_CARRIER_AIRLINE_ID', 'ORIGIN_AIRPORT_ID', 'DEST_AIRPORT_ID'])

Создание вложений работает должным образом, но я не могу получить petastorm чтобы вывести его, когда я смотрю на образец данных:

with converter_train.make_torch_dataloader(transform_spec=get_transform_spec_train()) as dataloader_train:
  sample = next(iter(dataloader_train))
  print(sample)

Я хочу иметь возможность использовать свои вложения, то есть shape (231440, 4), вместе с моими непрерывными числовыми функциями.