Reading ArrayRecord Files

Reading ArrayRecord Files#

Open in Colab

This tutorial provides an example of how to retrieve records from ArrayRecord files using grain.sources.ArrayRecordDataSource, also covers how to process and transform the data with Grain.

Install and Load Dependencies#

!pip install grain array_record

import pickle
import grain
import tensorflow_datasets as tfds
from array_record.python import array_record_module

Write a temp ArrayRecord file#

# Load a public tensorflow dataset.
test_tfds = tfds.data_source("bool_q", split="train")

# Write the dataset into a test array_record file.
example_file_path = "./test.array_record"
writer = array_record_module.ArrayRecordWriter(
    example_file_path, "group_size:1"
)
record_count = 0
for record in test_tfds:
  writer.write(pickle.dumps(record))
  record_count += 1
writer.close()

print(
    f"Number of records written to array_record file {example_file_path} :"
    f" {record_count}"
)

# @title Load Data Source
example_array_record_data_source = (grain.sources.ArrayRecordDataSource(
    example_file_path
))
print(f"Number of records: {len(example_array_record_data_source)}")

print(example_array_record_data_source[0])

Optimization Options#

Often times, ArrayRecord files can be heavily sharded on the order of thousands of files for just one dataset. These configurations can lead to a large number of readers and indices being kept in memory that can quickly use up memory in training pipelines. To specifically optimize ArrayRecord to conserve memory during the data processing stage, use the new reader_options to set the index_storage_option to offloaded. See the example below.

new_reader_options = {
    "index_storage_option": "offloaded"
}
datasource = grain.ArrayRecordDataSource(
    example_file_path,
    reader_options=new_reader_options
)

This option will offload readers indices when not in use and will prioritize memory efficiency over performance. On the other hand, the default index_storage_option is in_memory which will prioritize speed over memory.

Define Transformation Function#

# Load a pre trained tokenizer
from tokenizers import Tokenizer

tokenizer = Tokenizer.from_pretrained("bert-base-cased")

class ParseAndTokenizeText(grain.transforms.Map):
  """This function takes a serialized dict (as bytes), decodes it,

  applies a tokenizer to a specified feature within the dict,
  and returns the first 10 tokens from results.
  """

  def __init__(self, tokenizer, feature_name):
    self._tokenizer = tokenizer
    self._feature_name = feature_name

  def map(self, element: bytes) -> [str]:
    parsed_element = pickle.loads(element)
    # only pick the first 10 token IDs from the tokenized text for testing
    return self._tokenizer.encode(
        parsed_element[self._feature_name].decode('utf-8')
    ).tokens[:10]

Load and process data via the Dataset API#

# Example using Grain's MapDataset with ArrayRecord file source.
example_datasets = (
    grain.MapDataset.source(example_array_record_data_source)
    .shuffle(seed=42)
    .map(ParseAndTokenizeText(tokenizer, "question"))
    .batch(batch_size=10)
)

# Output a record at a random index
print(example_datasets[100])
Contents