47 lines · 1.5 KB · Python
RawDownload
"""Common utilities for the datasets""" import requestsfrom tqdm import tqdmimport numpy as np  def download_file(url: str, fname: str, chunk_size=1024):    """Helper function to download a file from a given url"""    resp = requests.get(url, stream=True)    total = int(resp.headers.get("content-length", 0))    with open(fname, "wb") as file, tqdm(        desc=fname,        total=total,        unit="iB",        unit_scale=True,        unit_divisor=1024,    ) as bar:        for data in resp.iter_content(chunk_size=chunk_size):            size = file.write(data)            bar.update(size)  def write_datafile(filename, toks):    """    Saves token data as a .bin file, for reading in C.    - First comes a header with 256 int32s    - The tokens follow, each as a uint16    """    assert len(toks) < 2**31, "token count too large" # ~2.1B tokens    # construct the header    header = np.zeros(256, dtype=np.int32)    header[0] = 20240520 # magic    header[1] = 1 # version    header[2] = len(toks) # number of tokens after the 256*4 bytes of header (each 2 bytes as uint16)    # validate that no token exceeds a uint16    maxtok = 2**16    assert all(0 <= t < maxtok for t in toks), "token dictionary too large for uint16"    # construct the tokens numpy array    toks_np = np.array(toks, dtype=np.uint16)    # write to file    print(f"writing {len(toks):,} tokens to {filename}")    with open(filename, "wb") as f:        f.write(header.tobytes())        f.write(toks_np.tobytes())