chore: import upstream snapshot with attribution

dev/data/data_common.py

@@ -0,0 +1,121 @@
+"""
+Common utilities for the datasets
+"""
+
+import requests
+from tqdm import tqdm
+import 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)
+
+
+HEADERS_INFO = {
+    "gpt-2": {
+        "magic": 20240520,
+        "version": 1,
+        "token_dtype": np.uint16,
+    },
+    "llama-3": {
+        "magic": 20240801,
+        "version": 7,
+        "token_dtype": np.uint32,
+    },
+}
+
+def write_datafile(filename, toks, model_desc="gpt-2"):
+    """
+    Saves token data as a .bin file, for reading in C.
+    - First comes a header with 256 int32s
+    - The tokens follow, each as uint16 (gpt-2) or uint32 (llama)
+    """
+    assert len(toks) < 2**31, "token count too large" # ~2.1B tokens
+    assert model_desc in ["gpt-2", "llama-3"], f"unknown model descriptor {model_desc}"
+    info = HEADERS_INFO[model_desc]
+    # construct the header
+    header = np.zeros(256, dtype=np.int32) # header is always 256 int32 values
+    header[0] = info["magic"]
+    header[1] = info["version"]
+    header[2] = len(toks) # number of tokens after the 256*4 bytes of header
+    # construct the data (numpy array of tokens)
+    toks_np = np.array(toks, dtype=info["token_dtype"])
+    # write to file
+    num_bytes = (256 * 4) + (len(toks) * toks_np.itemsize)
+    print(f"writing {len(toks):,} tokens to {filename} ({num_bytes:,} bytes) in the {model_desc} format")
+    with open(filename, "wb") as f:
+        f.write(header.tobytes())
+        f.write(toks_np.tobytes())
+
+def write_evalfile(filename, datas):
+    """
+    Saves eval data as a .bin file, for reading in C.
+    Used for multiple-choice style evals, e.g. HellaSwag and MMLU
+    - First comes a header with 256 int32s
+    - The examples follow, each example is a stream of uint16_t:
+        - <START_EXAMPLE> delimiter of 2**16-1, i.e. 65,535
+        - <EXAMPLE_BYTES>, bytes encoding this example, allowing efficient skip to next
+        - <EXAMPLE_INDEX>, the index of the example in the dataset
+        - <LABEL>, the index of the correct completion
+        - <NUM_COMPLETIONS>, indicating the number of completions (usually 4)
+        - <NUM><CONTEXT_TOKENS>, where <NUM> is the number of tokens in the context
+        - <NUM><COMPLETION_TOKENS>, repeated NUM_COMPLETIONS times
+    """
+    # construct the header
+    header = np.zeros(256, dtype=np.int32)
+    header[0] = 20240522 # magic
+    header[1] = 1 # version
+    header[2] = len(datas) # number of examples
+    header[3] = 0 # reserved for longest_example_bytes, fill in later
+    # now write the individual examples
+    longest_example_bytes = 0 # in units of uint16s
+    full_stream = [] # the stream of uint16s, we'll write a single time at the end
+    assert len(datas) < 2**16, "too many examples?"
+    for idx, data in enumerate(datas):
+        stream = []
+        # header of the example
+        stream.append(2**16-1) # <START_EXAMPLE>
+        stream.append(0) # <EXAMPLE_BYTES> (fill in later)
+        stream.append(idx) # <EXAMPLE_INDEX>
+        stream.append(data["label"]) # <LABEL>
+        ending_tokens = data["ending_tokens"]
+        assert len(ending_tokens) == 4, "expected 4 completions for now? can relax later"
+        stream.append(len(ending_tokens)) # <NUM_COMPLETIONS>
+        # the (shared) context tokens
+        ctx_tokens = data["ctx_tokens"]
+        assert all(0 <= t < 2**16-1 for t in ctx_tokens), "bad context token"
+        stream.append(len(ctx_tokens))
+        stream.extend(ctx_tokens)
+        # the completion tokens
+        for end_tokens in ending_tokens:
+            assert all(0 <= t < 2**16-1 for t in end_tokens), "bad completion token"
+            stream.append(len(end_tokens))
+            stream.extend(end_tokens)
+        # write to full stream
+        nbytes = len(stream)*2 # 2 bytes per uint16
+        assert nbytes < 2**16, "example too large?"
+        stream[1] = nbytes # fill in the <EXAMPLE_BYTES> field
+        longest_example_bytes = max(longest_example_bytes, nbytes)
+        full_stream.extend(stream)
+    # construct the numpy array
+    stream_np = np.array(full_stream, dtype=np.uint16)
+    # fill in the longest_example field
+    assert 0 < longest_example_bytes < 2**16, f"bad longest_example"
+    header[3] = longest_example_bytes
+    # write to file (for HellaSwag val this is 10,042 examples, 3.6MB file)
+    print(f"writing {len(datas):,} examples to {filename}")
+    with open(filename, "wb") as f:
+        f.write(header.tobytes())
+        f.write(stream_np.tobytes())