Как создать объект make_batch_reader библиотеки petastorm в DataBricks?

Question

У меня данные сохранены в паркетном формате. Petastorm — это библиотека, которую я использую для получения пакетов данных для обучения.

Хотя я смог сделать это в своей локальной системе, тот же код не работает в Databricks.

Код, который я использовал в своей локальной системе

      # create a iterator object train_reader. num_epochs is the number of epochs for which we want to train our model

with make_batch_reader('file:///config/workspace/scaled.parquet', num_epochs=4,shuffle_row_groups=False) as train_reader:
  train_ds = make_petastorm_dataset(train_reader).unbatch().map(lambda x: (tf.convert_to_tensor(x))).batch(2)
  

  for ele in train_ds:
    tensor = tf.reshape(ele,(2,1,15))
    model.fit(tensor,tensor)

Код, который я использовал в Databricks

      with make_batch_reader('dbfs://output/scaled.parquet', num_epochs=4,shuffle_row_groups=False) as train_reader:
    train_ds = make_petastorm_dataset(train_reader).unbatch().map(lambda x: (tf.convert_to_tensor(x))).batch(2)
  

    for ele in train_ds:
        tensor = tf.reshape(ele,(2,1,15))
        model.fit(tensor,tensor)

Ошибка, которую я получаю в коде DataBricks:

TypeError: init() отсутствуют 2 обязательных позиционных аргумента: «экземпляр» и «токен»

Я проверил документацию, но не смог найти ни одного аргумента, который идет по имени экземпляра и токена. Однако в аналогичном методе make_reader в petastorm для Azure Databricks я вижу следующую строку кода:

      # create sas token for storage account access, use your own adls account info
remote_url = "abfs://container_name@storage_account_url"
account_name = "<<adls account name>>"
linked_service_name = '<<linked service name>>'
TokenLibrary = spark._jvm.com.microsoft.azure.synapse.tokenlibrary.TokenLibrary
sas_token = TokenLibrary.getConnectionString(linked_service_name)

with make_reader('{}/data_directory'.format(remote_url), storage_options = {'sas_token' : sas_token}) as reader:
    for row in reader:
        print(row)

Здесь я вижу, что в качестве входных данных передается некоторый «sas_token».

Пожалуйста, предложите, как мне решить эту ошибку?

Я пробовал менять пути файла паркета, но у меня это не сработало.