OpenAI Week 15-16 // The Pile & Azure Blob Storage

For my project, I'm using The Pile, an 800 GB unstructured text dataset scraped from the web, representing a diverse range of modalities from medical papers to code. Such a large dataset is best stored on the cloud. I've written out a guide for how to get set up with The Pile hosted on Azure.

First, create an Azure Blob Storage account. I just followed the instructions and used the default configurations. Meanwhile, on your local or virtual machine, install the az CLI, blobfile, and boostedblob. This will help you easily interact with the files in your Azure storage account.

Now sign in with

 

az login

and open the link in a web browser to complete the sign in.

When you try to read one of your files from blob storage using blobfile, you'll get an error that you need to set the following environment variables first: AZURE_STORAGE_KEY, AZURE_CLIENT_ID, AZURE_CLIENT_SECRET, and AZURE_TENANT_ID. You can access these by typing

az ad sp create-for-rbac --name <your name> -o table

which will return a table with items like AppId, DisplayName, Name, Password, and Tenant. It's a bit opaque, but the mapping is

AppID $\rightarrow$ <AZURE_CLIENT_ID>

Name $\rightarrow$ <AZURE_STORAGE_KEY>

Password $\rightarrow$ <AZURE_CLIENT_SECRET>

Tenant $\rightarrow$ <AZURE_TENANT_ID>

so now you can set those environment variables:

import osos.environ['AZURE_STORAGE_KEY'] = <my storage key>
os.environ['AZURE_CLIENT_ID'] = <my client id>
os.environ['AZURE_CLIENT_SECRET'] = <my client secret>
os.environ['AZURE_TENANT_ID'] = <my tenant id>

And there you have it. Last but not least, I'm including my little script for reading off The Pile:

import zstandard
import jsonlines
import simdjson as json
import blobfile as bf
parser = json.Parser()

def json_parser(x):
  try:
    line = parser.parse(x).as_dict()
    return line
  except ValueError:
    return x
class PileReader:

  def __init__(self, filenames, para_joiner='\n\n'):
    if not isinstance(filenames, list):
      filenames = [filenames]
    self.filenames = filenames
    self.para_joiner = para_joiner

  def _read_fn(self, filename):
    with bf.BlobFile(filename, 'rb') as f:
      cctx = zstandard.ZstdDecompressor()
      reader_stream = io.BufferedReader(cctx.stream_reader(f))
      reader = jsonlines.Reader(reader_stream, loads=json_parser)
      for item in reader:
        result = dict()
        if isinstance(item, str):
          result['text'] = item
        else:
          text = item['text']
          if isinstance(text, list):
            text = self.para_joiner.join(text)
            result['text'] = text
        yield result

  def __iter__(self):
    for filename in self.filenames:
      return self._read_fn(filename)

Or you may wish to create a mini-file to work with first, by breaking off a piece of the first training file:

input_file = '/path/to/file/00.jsonl.zst'
output_file = '/path/to/file/00-mini.jsonl.zst'
line_counter = 0
num_lines_to_keep = 2000000   # ~4GB file

with bf.BlobFile(input_file, 'rb') as f_in:    with bf.BlobFile(output_file, 'wb') as f_out:        decompressor = zstandard.ZstdDecompressor()
        compressor = zstandard.ZstdCompressor()        reader_stream = io.BufferedReader(decompressor.stream_reader(f_in))
        reader = jsonlines.Reader(reader_stream, loads=json_parser)        writer_stream = io.BufferedWriter(compressor.stream_writer(f_out))
        writer = jsonlines.Writer(writer_stream)        for line in reader:            writer.write(line)            line_counter += 1
            if line_counter > num_lines_to_keep:
                break