IQ.Pilot Prebuilt Release @ ab07000

2026-08-22 23:42:42 -05:00
commit 9f9c9a70cc
3729 changed files with 778697 additions and 0 deletions
--- a/tinygrad_repo/extra/datasets/.gitignore
+++ b/tinygrad_repo/extra/datasets/.gitignore
@@ -0,0 +1,4 @@
+imagenet
+imagenet_bak
+mnist
+open-images-v6TEST
--- a/tinygrad_repo/extra/datasets/init.py
+++ b/tinygrad_repo/extra/datasets/init.py
@@ -0,0 +1,43 @@
+import os, gzip, tarfile, pickle
+import numpy as np
+from tinygrad import Tensor, dtypes
+from tinygrad.helpers import fetch
+
+def fetch_mnist(tensors=False):
+  parse = lambda file: np.frombuffer(gzip.open(file).read(), dtype=np.uint8).copy()
+  BASE_URL = "https://storage.googleapis.com/cvdf-datasets/mnist/"   # http://yann.lecun.com/exdb/mnist/ lacks https
+  X_train = parse(fetch(f"{BASE_URL}train-images-idx3-ubyte.gz"))[0x10:].reshape((-1, 28*28)).astype(np.float32)
+  Y_train = parse(fetch(f"{BASE_URL}train-labels-idx1-ubyte.gz"))[8:].astype(np.int8)
+  X_test = parse(fetch(f"{BASE_URL}t10k-images-idx3-ubyte.gz"))[0x10:].reshape((-1, 28*28)).astype(np.float32)
+  Y_test = parse(fetch(f"{BASE_URL}t10k-labels-idx1-ubyte.gz"))[8:].astype(np.int8)
+  if tensors: return Tensor(X_train).reshape(-1, 1, 28, 28), Tensor(Y_train), Tensor(X_test).reshape(-1, 1, 28, 28), Tensor(Y_test)
+  else: return X_train, Y_train, X_test, Y_test
+
+cifar_mean = [0.4913997551666284, 0.48215855929893703, 0.4465309133731618]
+cifar_std = [0.24703225141799082, 0.24348516474564, 0.26158783926049628]
+
+def fetch_cifar():
+  X_train = Tensor.empty(50000, 3*32*32, device=f'disk:/tmp/cifar_train_x', dtype=dtypes.uint8)
+  Y_train = Tensor.empty(50000, device=f'disk:/tmp/cifar_train_y', dtype=dtypes.int64)
+  X_test = Tensor.empty(10000, 3*32*32, device=f'disk:/tmp/cifar_test_x', dtype=dtypes.uint8)
+  Y_test = Tensor.empty(10000, device=f'disk:/tmp/cifar_test_y', dtype=dtypes.int64)
+
+  if not os.path.isfile("/tmp/cifar_extracted"):
+    def _load_disk_tensor(X, Y, db_list):
+      idx = 0
+      for db in db_list:
+        x, y = db[b'data'], np.array(db[b'labels'])
+        assert x.shape[0] == y.shape[0]
+        X[idx:idx+x.shape[0]].assign(x)
+        Y[idx:idx+x.shape[0]].assign(y)
+        idx += x.shape[0]
+      assert idx == X.shape[0] and X.shape[0] == Y.shape[0]
+
+    print("downloading and extracting CIFAR...")
+    fn = fetch('https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz')
+    tt = tarfile.open(fn, mode='r:gz')
+    _load_disk_tensor(X_train, Y_train, [pickle.load(tt.extractfile(f'cifar-10-batches-py/data_batch_{i}'), encoding="bytes") for i in range(1,6)])
+    _load_disk_tensor(X_test, Y_test, [pickle.load(tt.extractfile('cifar-10-batches-py/test_batch'), encoding="bytes")])
+    open("/tmp/cifar_extracted", "wb").close()
+
+  return X_train, Y_train, X_test, Y_test
--- a/tinygrad_repo/extra/datasets/fake_imagenet_from_mnist.py
+++ b/tinygrad_repo/extra/datasets/fake_imagenet_from_mnist.py
@@ -0,0 +1,41 @@
+#!/usr/bin/env python3
+import pathlib, json
+from tinygrad.helpers import trange
+from extra.datasets import fetch_mnist
+from PIL import Image
+import numpy as np
+from multiprocessing import Pool
+
+X_train, Y_train, X_test, Y_test = fetch_mnist()
+
+def act(arg):
+  (basedir, i, train) = arg
+  if train:
+    img = np.uint8(X_train[i]).reshape(28, 28)
+    nm = f"train/{Y_train[i]}/{i}.jpg"
+  else:
+    img = np.uint8(X_test[i]).reshape(28, 28)
+    nm = f"val/{Y_test[i]}/{i}.jpg"
+  Image.fromarray(img).resize((224, 224)).convert('RGB').save(basedir / nm)
+
+def create_fake_mnist_imagenet(basedir:pathlib.Path):
+  print(f"creating mock MNIST dataset at {basedir}")
+  basedir.mkdir(exist_ok=True)
+
+  with (basedir / "imagenet_class_index.json").open('w') as f:
+    f.write(json.dumps({str(i):[str(i), str(i)] for i in range(10)}))
+
+  for i in range(10):
+    (basedir / f"train/{i}").mkdir(parents=True, exist_ok=True)
+    (basedir / f"val/{i}").mkdir(parents=True, exist_ok=True)
+
+  def gen(train):
+    for idx in trange(X_train.shape[0] if train else X_test.shape[0]):
+      yield (basedir, idx, train)
+
+  with Pool(64) as p:
+    for _ in p.imap_unordered(act, gen(True)): pass
+    for _ in p.imap_unordered(act, gen(False)): pass
+
+if __name__ == "__main__":
+  create_fake_mnist_imagenet(pathlib.Path("./mnist"))
--- a/tinygrad_repo/extra/datasets/imagenet.py
+++ b/tinygrad_repo/extra/datasets/imagenet.py
@@ -0,0 +1,91 @@
+# for imagenet download prepare.sh and run it
+import glob, random, json, math
+import numpy as np
+from PIL import Image
+import functools, pathlib
+from tinygrad.helpers import diskcache, getenv
+
+@functools.cache
+def get_imagenet_categories():
+  ci = json.load(open(BASEDIR / "imagenet_class_index.json"))
+  return {v[0]: int(k) for k,v in ci.items()}
+
+if getenv("MNISTMOCK"):
+  BASEDIR = pathlib.Path(__file__).parent / "mnist"
+
+  @functools.cache
+  def get_train_files():
+    if not BASEDIR.exists():
+      from extra.datasets.fake_imagenet_from_mnist import create_fake_mnist_imagenet
+      create_fake_mnist_imagenet(BASEDIR)
+
+    if not (files:=glob.glob(p:=str(BASEDIR / "train/*/*"))): raise FileNotFoundError(f"No training files in {p}")
+    return files
+else:
+  BASEDIR = pathlib.Path(__file__).parent / "imagenet"
+
+  @diskcache
+  def get_train_files():
+    if not (files:=glob.glob(p:=str(BASEDIR / "train/*/*"))): raise FileNotFoundError(f"No training files in {p}")
+    return files
+
+@functools.cache
+def get_val_files():
+  if not (files:=glob.glob(p:=str(BASEDIR / "val/*/*"))): raise FileNotFoundError(f"No validation files in {p}")
+  return files
+
+def image_resize(img, size, interpolation):
+  w, h = img.size
+  w_new = int((w / h) * size) if w > h else size
+  h_new = int((h / w) * size) if h > w else size
+  return img.resize([w_new, h_new], interpolation)
+
+def rand_flip(img):
+  if random.random() < 0.5:
+    img = np.flip(img, axis=1).copy()
+  return img
+
+def center_crop(img):
+  rescale = min(img.size) / 256
+  crop_left = (img.width - 224 * rescale) / 2.0
+  crop_top = (img.height - 224 * rescale) / 2.0
+  img = img.resize((224, 224), Image.BILINEAR, box=(crop_left, crop_top, crop_left + 224 * rescale, crop_top + 224 * rescale))
+  return img
+
+# we don't use supplied imagenet bounding boxes, so scale min is just min_object_covered
+# https://github.com/tensorflow/tensorflow/blob/e193d8ea7776ef5c6f5d769b6fb9c070213e737a/tensorflow/core/kernels/image/sample_distorted_bounding_box_op.cc
+def random_resized_crop(img, size, scale=(0.10, 1.0), ratio=(3/4, 4/3)):
+  w, h = img.size
+  area = w * h
+
+  # Crop
+  random_solution_found = False
+  for _ in range(100):
+    aspect_ratio = random.uniform(ratio[0], ratio[1])
+    max_scale = min(min(w * aspect_ratio / h, h / aspect_ratio / w), scale[1])
+    target_area = area * random.uniform(scale[0], max_scale)
+
+    w_new = int(round(math.sqrt(target_area * aspect_ratio)))
+    h_new = int(round(math.sqrt(target_area / aspect_ratio)))
+
+    if 0 < w_new <= w and 0 < h_new <= h:
+      crop_left = random.randint(0, w - w_new)
+      crop_top = random.randint(0, h - h_new)
+
+      img = img.crop((crop_left, crop_top, crop_left + w_new, crop_top + h_new))
+      random_solution_found = True
+      break
+
+  if not random_solution_found:
+    # Center crop
+    img = center_crop(img)
+  else:
+    # Resize
+    img = img.resize([size, size], Image.BILINEAR)
+
+  return img
+
+def preprocess_train(img):
+  img = random_resized_crop(img, 224)
+  img = rand_flip(np.array(img))
+  return img
--- a/tinygrad_repo/extra/datasets/imagenet_download.py
+++ b/tinygrad_repo/extra/datasets/imagenet_download.py
@@ -0,0 +1,51 @@
+# Python version of https://gist.github.com/antoinebrl/7d00d5cb6c95ef194c737392ef7e476a
+from tinygrad.helpers import fetch
+from pathlib import Path
+from tqdm import tqdm
+import tarfile, os
+
+def imagenet_extract(file, path, small=False):
+  with tarfile.open(name=file) as tar:
+    if small: # Show progressbar only for big files
+      for member in tar.getmembers(): tar.extract(path=path, member=member)
+    else:
+      for member in tqdm(iterable=tar.getmembers(), total=len(tar.getmembers())): tar.extract(path=path, member=member)
+    tar.close()
+
+def imagenet_prepare_val():
+  # Read in the labels file
+  with open(Path(__file__).parent / "imagenet" / "imagenet_2012_validation_synset_labels.txt", 'r') as f:
+    labels = f.read().splitlines()
+  f.close()
+  # Get a list of images
+  images = os.listdir(Path(__file__).parent / "imagenet" / "val")
+  images.sort()
+  # Create folders and move files into those
+  for co,dir in enumerate(labels):
+    os.makedirs(Path(__file__).parent / "imagenet" / "val" / dir, exist_ok=True)
+    os.replace(Path(__file__).parent / "imagenet" / "val" / images[co], Path(__file__).parent / "imagenet" / "val" / dir / images[co])
+  os.remove(Path(__file__).parent / "imagenet" / "imagenet_2012_validation_synset_labels.txt")
+
+def imagenet_prepare_train():
+  images = os.listdir(Path(__file__).parent / "imagenet" / "train")
+  for co,tarf in enumerate(images):
+    # for each tar file found. Create a folder with its name. Extract into that folder. Remove tar file
+    if Path(Path(__file__).parent / "imagenet" / "train" / images[co]).is_file():
+      images[co] = tarf[:-4] # remove .tar from extracted tar files
+      os.makedirs(Path(__file__).parent / "imagenet" / "train" / images[co], exist_ok=True)
+      imagenet_extract(Path(__file__).parent / "imagenet" / "train" / tarf, Path(__file__).parent/ "imagenet" / "train" / images[co], small=True)
+      os.remove(Path(__file__).parent / "imagenet" / "train" / tarf)
+
+if __name__ == "__main__":
+  os.makedirs(Path(__file__).parent / "imagenet", exist_ok=True)
+  os.makedirs(Path(__file__).parent / "imagenet" / "val", exist_ok=True)
+  os.makedirs(Path(__file__).parent / "imagenet" / "train", exist_ok=True)
+  fetch("https://raw.githubusercontent.com/raghakot/keras-vis/master/resources/imagenet_class_index.json", Path(__file__).parent / "imagenet" / "imagenet_class_index.json")
+  fetch("https://raw.githubusercontent.com/tensorflow/models/master/research/slim/datasets/imagenet_2012_validation_synset_labels.txt", Path(__file__).parent / "imagenet"/ "imagenet_2012_validation_synset_labels.txt")
+  fetch("https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_val.tar", Path(__file__).parent / "imagenet" / "ILSVRC2012_img_val.tar") # 7GB
+  imagenet_extract(Path(__file__).parent / "imagenet" / "ILSVRC2012_img_val.tar", Path(__file__).parent / "imagenet" / "val")
+  imagenet_prepare_val()
+  if os.getenv('IMGNET_TRAIN', None) is not None:
+    fetch("https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_train.tar", Path(__file__).parent / "imagenet" / "ILSVRC2012_img_train.tar") #138GB!
+    imagenet_extract(Path(__file__).parent / "imagenet" / "ILSVRC2012_img_train.tar", Path(__file__).parent / "imagenet" / "train")
+    imagenet_prepare_train()
--- a/tinygrad_repo/extra/datasets/kits19.py
+++ b/tinygrad_repo/extra/datasets/kits19.py
@@ -0,0 +1,219 @@
+import random
+import functools
+from pathlib import Path
+import numpy as np
+import nibabel as nib
+from scipy import signal, ndimage
+import os
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import fetch
+
+BASEDIR = Path(__file__).parent / "kits19" / "data"
+TRAIN_PREPROCESSED_DIR =  Path(__file__).parent / "kits19" / "preprocessed" / "train"
+VAL_PREPROCESSED_DIR =  Path(__file__).parent / "kits19" / "preprocessed" / "val"
+
+@functools.cache
+def get_train_files():
+  return sorted([x for x in BASEDIR.iterdir() if x.stem.startswith("case") and int(x.stem.split("_")[-1]) < 210 and x not in get_val_files()])
+
+@functools.cache
+def get_val_files():
+  data = fetch("https://raw.githubusercontent.com/mlcommons/training/master/retired_benchmarks/unet3d/pytorch/evaluation_cases.txt").read_text()
+  return sorted([x for x in BASEDIR.iterdir() if x.stem.split("_")[-1] in data.split("\n")])
+
+def load_pair(file_path):
+  image, label = nib.load(file_path / "imaging.nii.gz"), nib.load(file_path / "segmentation.nii.gz")
+  image_spacings = image.header["pixdim"][1:4].tolist()
+  image, label = image.get_fdata().astype(np.float32), label.get_fdata().astype(np.uint8)
+  image, label = np.expand_dims(image, 0), np.expand_dims(label, 0)
+  return image, label, image_spacings
+
+def resample3d(image, label, image_spacings, target_spacing=(1.6, 1.2, 1.2)):
+  if image_spacings != target_spacing:
+    spc_arr, targ_arr, shp_arr = np.array(image_spacings), np.array(target_spacing), np.array(image.shape[1:])
+    new_shape = (spc_arr / targ_arr * shp_arr).astype(int).tolist()
+    image = F.interpolate(torch.from_numpy(np.expand_dims(image, axis=0)), size=new_shape, mode="trilinear", align_corners=True)
+    label = F.interpolate(torch.from_numpy(np.expand_dims(label, axis=0)), size=new_shape, mode="nearest")
+    image = np.squeeze(image.numpy(), axis=0)
+    label = np.squeeze(label.numpy(), axis=0)
+  return image, label
+
+def normal_intensity(image, min_clip=-79.0, max_clip=304.0, mean=101.0, std=76.9):
+  image = np.clip(image, min_clip, max_clip)
+  image = (image - mean) / std
+  return image
+
+def pad_to_min_shape(image, label, roi_shape=(128, 128, 128)):
+  current_shape = image.shape[1:]
+  bounds = [max(0, roi_shape[i] - current_shape[i]) for i in range(3)]
+  paddings = [(0, 0)] + [(bounds[i] // 2, bounds[i] - bounds[i] // 2) for i in range(3)]
+  image = np.pad(image, paddings, mode="edge")
+  label = np.pad(label, paddings, mode="edge")
+  return image, label
+
+def preprocess(file_path):
+  image, label, image_spacings = load_pair(file_path)
+  image, label = resample3d(image, label, image_spacings)
+  image = normal_intensity(image.copy())
+  image, label = pad_to_min_shape(image, label)
+  return image, label
+
+def preprocess_dataset(filenames, preprocessed_dir, val):
+  if not preprocessed_dir.is_dir(): os.makedirs(preprocessed_dir)
+  for fn in tqdm(filenames, desc=f"preprocessing {'validation' if val else 'training'}"):
+    case = os.path.basename(fn)
+    image, label = preprocess(fn)
+    image, label = image.astype(np.float32), label.astype(np.uint8)
+    np.save(preprocessed_dir / f"{case}_x.npy", image, allow_pickle=False)
+    np.save(preprocessed_dir / f"{case}_y.npy", label, allow_pickle=False)
+
+def iterate(files, preprocessed_dir=None, val=True, shuffle=False, bs=1):
+  order = list(range(0, len(files)))
+  if shuffle: random.shuffle(order)
+  for i in range(0, len(files), bs):
+    samples = []
+    for i in order[i:i+bs]:
+      if preprocessed_dir is not None:
+        x_cached_path, y_cached_path = preprocessed_dir / f"{os.path.basename(files[i])}_x.npy", preprocessed_dir / f"{os.path.basename(files[i])}_y.npy"
+        if x_cached_path.exists() and y_cached_path.exists():
+          samples += [(np.load(x_cached_path), np.load(y_cached_path))]
+      else: samples += [preprocess(files[i])]
+    X, Y = [x[0] for x in samples], [x[1] for x in samples]
+    if val:
+      yield X[0][None], Y[0]
+    else:
+      X_preprocessed, Y_preprocessed = [], []
+      for x, y in zip(X, Y):
+        x, y = rand_balanced_crop(x, y)
+        x, y = rand_flip(x, y)
+        x, y = x.astype(np.float32), y.astype(np.uint8)
+        x = random_brightness_augmentation(x)
+        x = gaussian_noise(x)
+        X_preprocessed.append(x)
+        Y_preprocessed.append(y)
+      yield np.stack(X_preprocessed, axis=0), np.stack(Y_preprocessed, axis=0)
+
+def gaussian_kernel(n, std):
+  gaussian_1d = signal.windows.gaussian(n, std)
+  gaussian_2d = np.outer(gaussian_1d, gaussian_1d)
+  gaussian_3d = np.outer(gaussian_2d, gaussian_1d)
+  gaussian_3d = gaussian_3d.reshape(n, n, n)
+  gaussian_3d = np.cbrt(gaussian_3d)
+  gaussian_3d /= gaussian_3d.max()
+  return gaussian_3d
+
+def pad_input(volume, roi_shape, strides, padding_mode="constant", padding_val=-2.2, dim=3):
+  bounds = [(strides[i] - volume.shape[2:][i] % strides[i]) % strides[i] for i in range(dim)]
+  bounds = [bounds[i] if (volume.shape[2:][i] + bounds[i]) >= roi_shape[i] else bounds[i] + strides[i] for i in range(dim)]
+  paddings = [bounds[2]//2, bounds[2]-bounds[2]//2, bounds[1]//2, bounds[1]-bounds[1]//2, bounds[0]//2, bounds[0]-bounds[0]//2, 0, 0, 0, 0]
+  return F.pad(torch.from_numpy(volume), paddings, mode=padding_mode, value=padding_val).numpy(), paddings
+
+def sliding_window_inference(model, inputs, labels, roi_shape=(128, 128, 128), overlap=0.5, gpus=None):
+  from tinygrad.engine.jit import TinyJit
+  mdl_run = TinyJit(lambda x: model(x).realize())
+  image_shape, dim = list(inputs.shape[2:]), len(inputs.shape[2:])
+  strides = [int(roi_shape[i] * (1 - overlap)) for i in range(dim)]
+  bounds = [image_shape[i] % strides[i] for i in range(dim)]
+  bounds = [bounds[i] if bounds[i] < strides[i] // 2 else 0 for i in range(dim)]
+  inputs = inputs[
+    ...,
+    bounds[0]//2:image_shape[0]-(bounds[0]-bounds[0]//2),
+    bounds[1]//2:image_shape[1]-(bounds[1]-bounds[1]//2),
+    bounds[2]//2:image_shape[2]-(bounds[2]-bounds[2]//2),
+  ]
+  labels = labels[
+    ...,
+    bounds[0]//2:image_shape[0]-(bounds[0]-bounds[0]//2),
+    bounds[1]//2:image_shape[1]-(bounds[1]-bounds[1]//2),
+    bounds[2]//2:image_shape[2]-(bounds[2]-bounds[2]//2),
+  ]
+  inputs, paddings = pad_input(inputs, roi_shape, strides)
+  padded_shape = inputs.shape[2:]
+  size = [(inputs.shape[2:][i] - roi_shape[i]) // strides[i] + 1 for i in range(dim)]
+  result = np.zeros((1, 3, *padded_shape), dtype=np.float32)
+  norm_map = np.zeros((1, 3, *padded_shape), dtype=np.float32)
+  norm_patch = gaussian_kernel(roi_shape[0], 0.125 * roi_shape[0])
+  norm_patch = np.expand_dims(norm_patch, axis=0)
+  for i in range(0, strides[0] * size[0], strides[0]):
+    for j in range(0, strides[1] * size[1], strides[1]):
+      for k in range(0, strides[2] * size[2], strides[2]):
+        out = mdl_run(Tensor(inputs[..., i:roi_shape[0]+i,j:roi_shape[1]+j, k:roi_shape[2]+k], device=gpus)).numpy()
+        result[..., i:roi_shape[0]+i, j:roi_shape[1]+j, k:roi_shape[2]+k] += out * norm_patch
+        norm_map[..., i:roi_shape[0]+i, j:roi_shape[1]+j, k:roi_shape[2]+k] += norm_patch
+  result /= norm_map
+  result = result[..., paddings[4]:image_shape[0]+paddings[4], paddings[2]:image_shape[1]+paddings[2], paddings[0]:image_shape[2]+paddings[0]]
+  return result, labels
+
+def rand_flip(image, label, axis=(1, 2, 3)):
+  prob = 1 / len(axis)
+  for ax in axis:
+    if random.random() < prob:
+      image = np.flip(image, axis=ax).copy()
+      label = np.flip(label, axis=ax).copy()
+  return image, label
+
+def random_brightness_augmentation(image, low=0.7, high=1.3, prob=0.1):
+  if random.random() < prob:
+    factor = np.random.uniform(low=low, high=high, size=1)
+    image = (image * (1 + factor)).astype(image.dtype)
+  return image
+
+def gaussian_noise(image, mean=0.0, std=0.1, prob=0.1):
+  if random.random() < prob:
+    scale = np.random.uniform(low=0.0, high=std)
+    noise = np.random.normal(loc=mean, scale=scale, size=image.shape).astype(image.dtype)
+    image += noise
+  return image
+
+def _rand_foreg_cropb(image, label, patch_size):
+  def adjust(foreg_slice, label, idx):
+    diff = patch_size[idx - 1] - (foreg_slice[idx].stop - foreg_slice[idx].start)
+    sign = -1 if diff < 0 else 1
+    diff = abs(diff)
+    ladj = 0 if diff == 0 else random.randrange(diff)
+    hadj = diff - ladj
+    low = max(0, foreg_slice[idx].start - sign * ladj)
+    high = min(label.shape[idx], foreg_slice[idx].stop + sign * hadj)
+    diff = patch_size[idx - 1] - (high - low)
+    if diff > 0 and low == 0: high += diff
+    elif diff > 0: low -= diff
+    return low, high
+
+  cl = np.random.choice(np.unique(label[label > 0]))
+  foreg_slices = ndimage.find_objects(ndimage.label(label==cl)[0])
+  foreg_slices = [x for x in foreg_slices if x is not None]
+  slice_volumes = [np.prod([s.stop - s.start for s in sl]) for sl in foreg_slices]
+  slice_idx = np.argsort(slice_volumes)[-2:]
+  foreg_slices = [foreg_slices[i] for i in slice_idx]
+  if not foreg_slices: return _rand_crop(image, label)
+  foreg_slice = foreg_slices[random.randrange(len(foreg_slices))]
+  low_x, high_x = adjust(foreg_slice, label, 1)
+  low_y, high_y = adjust(foreg_slice, label, 2)
+  low_z, high_z = adjust(foreg_slice, label, 3)
+  image = image[:, low_x:high_x, low_y:high_y, low_z:high_z]
+  label = label[:, low_x:high_x, low_y:high_y, low_z:high_z]
+  return image, label
+
+def _rand_crop(image, label, patch_size):
+  ranges = [s - p for s, p in zip(image.shape[1:], patch_size)]
+  cord = [0 if x == 0 else random.randrange(x) for x in ranges]
+  low_x, high_x = cord[0], cord[0] + patch_size[0]
+  low_y, high_y = cord[1], cord[1] + patch_size[1]
+  low_z, high_z = cord[2], cord[2] + patch_size[2]
+  image = image[:, low_x:high_x, low_y:high_y, low_z:high_z]
+  label = label[:, low_x:high_x, low_y:high_y, low_z:high_z]
+  return image, label
+
+def rand_balanced_crop(image, label, patch_size=(128, 128, 128), oversampling=0.4):
+  if random.random() < oversampling:
+    image, label = _rand_foreg_cropb(image, label, patch_size)
+  else:
+    image, label = _rand_crop(image, label, patch_size)
+  return image, label
+
+if __name__ == "__main__":
+  for X, Y in iterate(get_val_files()):
+    print(X.shape, Y.shape)
--- a/tinygrad_repo/extra/datasets/librispeech.py
+++ b/tinygrad_repo/extra/datasets/librispeech.py
@@ -0,0 +1,82 @@
+import json
+import pathlib
+import numpy as np
+import librosa
+import soundfile
+
+"""
+The dataset has to be downloaded manually from https://www.openslr.org/12/ and put in `extra/datasets/librispeech`.
+For mlperf validation the dev-clean dataset is used.
+
+Then all the flacs have to be converted to wav using something like:
+```fish
+for file in $(find * | grep flac); do ffmpeg -i $file -ar 16k "$(dirname $file)/$(basename $file .flac).wav"; done
+```
+
+Then this [file](https://github.com/mlcommons/inference/blob/master/speech_recognition/rnnt/dev-clean-wav.json) has to also be put in `extra/datasets/librispeech`.
+"""
+BASEDIR = pathlib.Path(__file__).parent / "librispeech"
+with open(BASEDIR / "dev-clean-wav.json") as f:
+  ci = json.load(f)
+
+FILTER_BANK = np.expand_dims(librosa.filters.mel(sr=16000, n_fft=512, n_mels=80, fmin=0, fmax=8000), 0)
+WINDOW = librosa.filters.get_window("hann", 320)
+
+def feature_extract(x, x_lens):
+  x_lens = np.ceil((x_lens / 160) / 3).astype(np.int32)
+
+  # pre-emphasis
+  x = np.concatenate((np.expand_dims(x[:, 0], 1), x[:, 1:] - 0.97 * x[:, :-1]), axis=1)
+
+  # stft
+  x = librosa.stft(x, n_fft=512, window=WINDOW, hop_length=160, win_length=320, center=True, pad_mode="reflect")
+  x = np.stack((x.real, x.imag), axis=-1)
+
+  # power spectrum
+  x = (x**2).sum(-1)
+
+  # mel filter bank
+  x = np.matmul(FILTER_BANK, x)
+
+  # log
+  x = np.log(x + 1e-20)
+
+  # feature splice
+  seq = [x]
+  for i in range(1, 3):
+    tmp = np.zeros_like(x)
+    tmp[:, :, :-i] = x[:, :, i:]
+    seq.append(tmp)
+  features = np.concatenate(seq, axis=1)[:, :, ::3]
+
+  # normalize
+  features_mean = np.zeros((features.shape[0], features.shape[1]), dtype=np.float32)
+  features_std = np.zeros((features.shape[0], features.shape[1]), dtype=np.float32)
+  for i in range(features.shape[0]):
+    features_mean[i, :] = features[i, :, :x_lens[i]].mean(axis=1)
+    features_std[i, :] = features[i, :, :x_lens[i]].std(axis=1, ddof=1)
+  features_std += 1e-5
+  features = (features - np.expand_dims(features_mean, 2)) / np.expand_dims(features_std, 2)
+
+  return features.transpose(2, 0, 1), x_lens.astype(np.float32)
+
+def load_wav(file):
+  sample = soundfile.read(file)[0].astype(np.float32)
+  return sample, sample.shape[0]
+
+def iterate(bs=1, start=0):
+  print(f"there are {len(ci)} samples in the dataset")
+  for i in range(start, len(ci), bs):
+    samples, sample_lens = zip(*[load_wav(BASEDIR / v["files"][0]["fname"]) for v in ci[i : i + bs]])
+    samples = list(samples)
+    # pad to same length
+    max_len = max(sample_lens)
+    for j in range(len(samples)):
+      samples[j] = np.pad(samples[j], (0, max_len - sample_lens[j]), "constant")
+    samples, sample_lens = np.array(samples), np.array(sample_lens)
+
+    yield feature_extract(samples, sample_lens), np.array([v["transcript"] for v in ci[i : i + bs]])
+
+if __name__ == "__main__":
+  X, Y = next(iterate())
+  print(X[0].shape, Y.shape)
--- a/tinygrad_repo/extra/datasets/openimages.py
+++ b/tinygrad_repo/extra/datasets/openimages.py
@@ -0,0 +1,209 @@
+import glob
+import sys
+import json
+import numpy as np
+from PIL import Image
+from pathlib import Path
+import boto3, botocore
+from tinygrad import Tensor, dtypes
+from tinygrad.helpers import fetch, tqdm, getenv
+import pandas as pd
+import concurrent.futures
+
+BASEDIR = Path(__file__).parent / "open-images-v6-mlperf"
+BUCKET_NAME = "open-images-dataset"
+TRAIN_BBOX_ANNOTATIONS_URL = "https://storage.googleapis.com/openimages/v6/oidv6-train-annotations-bbox.csv"
+VALIDATION_BBOX_ANNOTATIONS_URL = "https://storage.googleapis.com/openimages/v5/validation-annotations-bbox.csv"
+MAP_CLASSES_URL = "https://storage.googleapis.com/openimages/v5/class-descriptions-boxable.csv"
+MLPERF_CLASSES = ['Airplane', 'Antelope', 'Apple', 'Backpack', 'Balloon', 'Banana',
+  'Barrel', 'Baseball bat', 'Baseball glove', 'Bee', 'Beer', 'Bench', 'Bicycle',
+  'Bicycle helmet', 'Bicycle wheel', 'Billboard', 'Book', 'Bookcase', 'Boot',
+  'Bottle', 'Bowl', 'Bowling equipment', 'Box', 'Boy', 'Brassiere', 'Bread',
+  'Broccoli', 'Bronze sculpture', 'Bull', 'Bus', 'Bust', 'Butterfly', 'Cabinetry',
+  'Cake', 'Camel', 'Camera', 'Candle', 'Candy', 'Cannon', 'Canoe', 'Carrot', 'Cart',
+  'Castle', 'Cat', 'Cattle', 'Cello', 'Chair', 'Cheese', 'Chest of drawers', 'Chicken',
+  'Christmas tree', 'Coat', 'Cocktail', 'Coffee', 'Coffee cup', 'Coffee table', 'Coin',
+  'Common sunflower', 'Computer keyboard', 'Computer monitor', 'Convenience store',
+  'Cookie', 'Countertop', 'Cowboy hat', 'Crab', 'Crocodile', 'Cucumber', 'Cupboard',
+  'Curtain', 'Deer', 'Desk', 'Dinosaur', 'Dog', 'Doll', 'Dolphin', 'Door', 'Dragonfly',
+  'Drawer', 'Dress', 'Drum', 'Duck', 'Eagle', 'Earrings', 'Egg (Food)', 'Elephant',
+  'Falcon', 'Fedora', 'Flag', 'Flowerpot', 'Football', 'Football helmet', 'Fork',
+  'Fountain', 'French fries', 'French horn', 'Frog', 'Giraffe', 'Girl', 'Glasses',
+  'Goat', 'Goggles', 'Goldfish', 'Gondola', 'Goose', 'Grape', 'Grapefruit', 'Guitar',
+  'Hamburger', 'Handbag', 'Harbor seal', 'Headphones', 'Helicopter', 'High heels',
+  'Hiking equipment', 'Horse', 'House', 'Houseplant', 'Human arm', 'Human beard',
+  'Human body', 'Human ear', 'Human eye', 'Human face', 'Human foot', 'Human hair',
+  'Human hand', 'Human head', 'Human leg', 'Human mouth', 'Human nose', 'Ice cream',
+  'Jacket', 'Jeans', 'Jellyfish', 'Juice', 'Kitchen & dining room table', 'Kite',
+  'Lamp', 'Lantern', 'Laptop', 'Lavender (Plant)', 'Lemon', 'Light bulb', 'Lighthouse',
+  'Lily', 'Lion', 'Lipstick', 'Lizard', 'Man', 'Maple', 'Microphone', 'Mirror',
+  'Mixing bowl', 'Mobile phone', 'Monkey', 'Motorcycle', 'Muffin', 'Mug', 'Mule',
+  'Mushroom', 'Musical keyboard', 'Necklace', 'Nightstand', 'Office building',
+  'Orange', 'Owl', 'Oyster', 'Paddle', 'Palm tree', 'Parachute', 'Parrot', 'Pen',
+  'Penguin', 'Personal flotation device', 'Piano', 'Picture frame', 'Pig', 'Pillow',
+  'Pizza', 'Plate', 'Platter', 'Porch', 'Poster', 'Pumpkin', 'Rabbit', 'Rifle',
+  'Roller skates', 'Rose', 'Salad', 'Sandal', 'Saucer', 'Saxophone', 'Scarf', 'Sea lion',
+  'Sea turtle', 'Sheep', 'Shelf', 'Shirt', 'Shorts', 'Shrimp', 'Sink', 'Skateboard',
+  'Ski', 'Skull', 'Skyscraper', 'Snake', 'Sock', 'Sofa bed', 'Sparrow', 'Spider', 'Spoon',
+  'Sports uniform', 'Squirrel', 'Stairs', 'Stool', 'Strawberry', 'Street light',
+  'Studio couch', 'Suit', 'Sun hat', 'Sunglasses', 'Surfboard', 'Sushi', 'Swan',
+  'Swimming pool', 'Swimwear', 'Tank', 'Tap', 'Taxi', 'Tea', 'Teddy bear', 'Television',
+  'Tent', 'Tie', 'Tiger', 'Tin can', 'Tire', 'Toilet', 'Tomato', 'Tortoise', 'Tower',
+  'Traffic light', 'Train', 'Tripod', 'Truck', 'Trumpet', 'Umbrella', 'Van', 'Vase',
+  'Vehicle registration plate', 'Violin', 'Wall clock', 'Waste container', 'Watch',
+  'Whale', 'Wheel', 'Wheelchair', 'Whiteboard', 'Window', 'Wine', 'Wine glass', 'Woman',
+  'Zebra', 'Zucchini',
+]
+
+
+def openimages(base_dir:Path, subset:str, ann_file:Path):
+  valid_subsets = ['train', 'validation']
+  if subset not in valid_subsets:
+    raise ValueError(f"{subset=} must be one of {valid_subsets}")
+
+  fetch_openimages(ann_file, base_dir, subset)
+
+# this slows down the conversion a lot!
+# maybe use https://raw.githubusercontent.com/scardine/image_size/master/get_image_size.py
+def extract_dims(path): return Image.open(path).size[::-1]
+
+def export_to_coco(class_map, annotations, image_list, dataset_path, output_path, subset, classes=MLPERF_CLASSES):
+  output_path.parent.mkdir(parents=True, exist_ok=True)
+  cats = [{"id": i, "name": c, "supercategory": None} for i, c in enumerate(classes)]
+  categories_map = pd.DataFrame([(i, c) for i, c in enumerate(classes)], columns=["category_id", "category_name"])
+  class_map = class_map.merge(categories_map, left_on="DisplayName", right_on="category_name", how="inner")
+  annotations = annotations[annotations["ImageID"].isin(image_list)]
+  annotations = annotations.merge(class_map, on="LabelName", how="inner")
+  annotations["image_id"] = pd.factorize(annotations["ImageID"].tolist())[0]
+  annotations[["height", "width"]] = annotations.apply(lambda x: extract_dims(dataset_path / f"{x['ImageID']}.jpg"), axis=1, result_type="expand")
+
+  # Images
+  imgs = [{"id": int(id + 1), "file_name": f"{image_id}.jpg", "height": row["height"], "width": row["width"], "subset": subset, "license": None, "coco_url": None}
+    for (id, image_id), row in (annotations.groupby(["image_id", "ImageID"]).first().iterrows())
+  ]
+
+  # Annotations
+  annots = []
+  for i, row in annotations.iterrows():
+    xmin, ymin, xmax, ymax, img_w, img_h = [row[k] for k in ["XMin", "YMin", "XMax", "YMax", "width", "height"]]
+    x, y, w, h = xmin * img_w, ymin * img_h, (xmax - xmin) * img_w, (ymax - ymin) * img_h
+    coco_annot = {"id": int(i) + 1, "image_id": int(row["image_id"] + 1), "category_id": int(row["category_id"]), "bbox": [x, y, w, h], "area": w * h}
+    coco_annot.update({k: row[k] for k in ["IsOccluded", "IsInside", "IsDepiction", "IsTruncated", "IsGroupOf"]})
+    coco_annot["iscrowd"] = int(row["IsGroupOf"])
+    annots.append(coco_annot)
+
+  info = {"dataset": "openimages_mlperf", "version": "v6"}
+  coco_annotations = {"info": info, "licenses": [], "categories": cats, "images": imgs, "annotations": annots}
+  with open(output_path, "w") as fp:
+    json.dump(coco_annotations, fp)
+
+def get_image_list(class_map, annotations, classes=MLPERF_CLASSES):
+  labels = class_map[class_map["DisplayName"].isin(classes)]["LabelName"]
+  image_ids = annotations[annotations["LabelName"].isin(labels)]["ImageID"].unique()
+  return image_ids
+
+def download_image(bucket, subset, image_id, data_dir):
+  try:
+    bucket.download_file(f"{subset}/{image_id}.jpg", f"{data_dir}/{image_id}.jpg")
+  except botocore.exceptions.ClientError as exception:
+    sys.exit(f"ERROR when downloading image `validation/{image_id}`: {str(exception)}")
+
+def fetch_openimages(output_fn:str, base_dir:Path, subset:str):
+  bucket = boto3.resource("s3", config=botocore.config.Config(signature_version=botocore.UNSIGNED)).Bucket(BUCKET_NAME)
+
+  annotations_dir, data_dir = base_dir / "annotations", base_dir / f"{subset}/data"
+  annotations_dir.mkdir(parents=True, exist_ok=True)
+  data_dir.mkdir(parents=True, exist_ok=True)
+
+  if subset == "train":
+    annotations_fn = annotations_dir / TRAIN_BBOX_ANNOTATIONS_URL.split('/')[-1]
+    fetch(TRAIN_BBOX_ANNOTATIONS_URL, annotations_fn)
+  else:  # subset == validation
+    annotations_fn = annotations_dir / VALIDATION_BBOX_ANNOTATIONS_URL.split('/')[-1]
+    fetch(VALIDATION_BBOX_ANNOTATIONS_URL, annotations_fn)
+
+  annotations = pd.read_csv(annotations_fn)
+
+  classmap_fn = annotations_dir / MAP_CLASSES_URL.split('/')[-1]
+  fetch(MAP_CLASSES_URL, classmap_fn)
+  class_map = pd.read_csv(classmap_fn, names=["LabelName", "DisplayName"])
+
+  image_list = get_image_list(class_map, annotations)
+
+  with concurrent.futures.ThreadPoolExecutor() as executor:
+    futures = [executor.submit(download_image, bucket, subset, image_id, data_dir) for image_id in image_list]
+    for future in (t := tqdm(concurrent.futures.as_completed(futures), total=len(image_list))):
+      t.set_description(f"Downloading images")
+      future.result()
+
+  print("Converting annotations to COCO format...")
+  export_to_coco(class_map, annotations, image_list, data_dir, output_fn, subset)
+
+def image_load(base_dir, subset, fn):
+  img_folder = base_dir / f"{subset}/data"
+  return Image.open(img_folder / fn).convert('RGB')
+
+def prepare_target(annotations, img_id, img_size):
+  boxes = [annot["bbox"] for annot in annotations]
+  boxes = np.array(boxes, dtype=np.float32).reshape(-1, 4)
+  boxes[:, 2:] += boxes[:, :2]
+  boxes[:, 0::2] = boxes[:, 0::2].clip(0, img_size[1])
+  boxes[:, 1::2] = boxes[:, 1::2].clip(0, img_size[0])
+  keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0])
+  boxes = boxes[keep]
+  classes = [annot["category_id"] for annot in annotations]
+  classes = np.array(classes, dtype=np.int64)
+  classes = classes[keep]
+  return {"boxes": boxes, "labels": classes, "image_id": img_id, "image_size": img_size}
+
+def download_dataset(base_dir:Path, subset:str) -> Path:
+  if (ann_file:=base_dir / f"{subset}/labels/openimages-mlperf.json").is_file(): print(f"{subset} dataset is already available")
+  else:
+    print(f"Downloading {subset} dataset...")
+    openimages(base_dir, subset, ann_file)
+    print("Done")
+
+  return ann_file
+
+def random_horizontal_flip(img, tgt, prob=0.5):
+  import torch
+  import torchvision.transforms.functional as F
+  if torch.rand(1) < prob:
+    w = img.size[0]
+    img = F.hflip(img)
+    tgt["boxes"][:, [0, 2]] = w - tgt["boxes"][:, [2, 0]]
+  return img, tgt
+
+def resize(img:Image, tgt:dict[str, np.ndarray|tuple]|None=None, size:tuple[int, int]=(800, 800)) -> tuple[np.ndarray, np.ndarray, tuple]|tuple[np.ndarray, tuple]:
+  import torchvision.transforms.functional as F
+  img_size = img.size[::-1]
+  img = F.resize(img, size=size)
+  img = np.array(img)
+
+  if tgt is not None:
+    ratios = [s / s_orig for s, s_orig in zip(size, img_size)]
+    ratio_h, ratio_w = ratios
+    x_min, y_min, x_max, y_max = [tgt["boxes"][:, i] for i in range(tgt["boxes"].shape[-1])]
+    x_min = x_min * ratio_w
+    x_max = x_max * ratio_w
+    y_min = y_min * ratio_h
+    y_max = y_max * ratio_h
+
+    tgt["boxes"] = np.stack([x_min, y_min, x_max, y_max], axis=1)
+    return img, tgt, img_size
+
+  return img, img_size
+
+def normalize(img:Tensor, device:list[str]|None = None):
+  mean = Tensor([0.485, 0.456, 0.406], device=device, dtype=dtypes.float32).reshape(1, -1, 1, 1)
+  std = Tensor([0.229, 0.224, 0.225], device=device, dtype=dtypes.float32).reshape(1, -1, 1, 1)
+  img = ((img.permute([0, 3, 1, 2]) / 255.0) - mean) / std
+  return img.cast(dtypes.default_float)
+
+def get_dataset_count(base_dir:Path, val:bool) -> int:
+  if not (files:=glob.glob(p:=str(base_dir / f"{'validation' if val else 'train'}/data/*.jpg"))): raise FileNotFoundError(f"No files in {p}")
+  return len(files)
+
+if __name__ == "__main__":
+  download_dataset(base_dir:=getenv("BASEDIR", BASEDIR), "train")
+  download_dataset(base_dir, "validation")
--- a/tinygrad_repo/extra/datasets/preprocess_imagenet.py
+++ b/tinygrad_repo/extra/datasets/preprocess_imagenet.py
@@ -0,0 +1,21 @@
+from tinygrad import Tensor, dtypes
+from extra.datasets.imagenet import iterate, get_val_files
+
+if __name__ == "__main__":
+  #sz = len(get_val_files())
+  sz = 32*100
+  X,Y = None, None
+
+  idx = 0
+  for x,y in iterate(shuffle=False):
+    print(x.shape, y.shape, x.dtype, y.dtype)
+    assert x.shape[0] == y.shape[0]
+    bs = x.shape[0]
+    if X is None:
+      X = Tensor.empty(sz, *x.shape[1:], device="disk:/tmp/imagenet_x", dtype=dtypes.uint8)
+      Y = Tensor.empty(sz, *y.shape[1:], device="disk:/tmp/imagenet_y", dtype=dtypes.int64)
+      print(X.shape, Y.shape)
+    X[idx:idx+bs].assign(x)
+    Y[idx:idx+bs].assign(y)
+    idx += bs
+    if idx >= sz: break
--- a/tinygrad_repo/extra/datasets/sops.gz
+++ b/tinygrad_repo/extra/datasets/sops.gz
--- a/tinygrad_repo/extra/datasets/squad.py
+++ b/tinygrad_repo/extra/datasets/squad.py
@@ -0,0 +1,148 @@
+import json
+import os
+from pathlib import Path
+from transformers import BertTokenizer
+import numpy as np
+from tinygrad.helpers import fetch
+
+BASEDIR = Path(__file__).parent / "squad"
+def init_dataset():
+  os.makedirs(BASEDIR, exist_ok=True)
+  fetch("https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v1.1.json", BASEDIR / "dev-v1.1.json")
+  with open(BASEDIR / "dev-v1.1.json") as f:
+    data = json.load(f)["data"]
+
+  examples = []
+  for article in data:
+    for paragraph in article["paragraphs"]:
+      text = paragraph["context"]
+      doc_tokens = []
+      prev_is_whitespace = True
+      for c in text:
+        if c == " " or c == "\t" or c == "\r" or c == "\n" or ord(c) == 0x202F:
+          prev_is_whitespace = True
+        else:
+          if prev_is_whitespace:
+            doc_tokens.append(c)
+          else:
+            doc_tokens[-1] += c
+          prev_is_whitespace = False
+
+      for qa in paragraph["qas"]:
+        qa_id = qa["id"]
+        q_text = qa["question"]
+
+        examples.append({
+          "id": qa_id,
+          "question": q_text,
+          "context": doc_tokens,
+          "answers": list(map(lambda x: x["text"], qa["answers"]))
+        })
+  return examples
+
+def _check_is_max_context(doc_spans, cur_span_index, position):
+  best_score, best_span_index = None, None
+  for di, (doc_start, doc_length) in enumerate(doc_spans):
+    end = doc_start + doc_length - 1
+    if position < doc_start:
+      continue
+    if position > end:
+      continue
+    num_left_context = position - doc_start
+    num_right_context = end - position
+    score = min(num_left_context, num_right_context) + 0.01 * doc_length
+    if best_score is None or score > best_score:
+      best_score = score
+      best_span_index = di
+  return cur_span_index == best_span_index
+
+def convert_example_to_features(example, tokenizer):
+  query_tokens = tokenizer.tokenize(example["question"])
+
+  if len(query_tokens) > 64:
+    query_tokens = query_tokens[:64]
+
+  tok_to_orig_index = []
+  orig_to_tok_index = []
+  all_doc_tokens = []
+  for i, token in enumerate(example["context"]):
+    orig_to_tok_index.append(len(all_doc_tokens))
+    sub_tokens = tokenizer.tokenize(token)
+    for sub_token in sub_tokens:
+      tok_to_orig_index.append(i)
+      all_doc_tokens.append(sub_token)
+
+  max_tokens_for_doc = 384 - len(query_tokens) - 3
+
+  doc_spans = []
+  start_offset = 0
+  while start_offset < len(all_doc_tokens):
+    length = len(all_doc_tokens) - start_offset
+    length = min(length, max_tokens_for_doc)
+    doc_spans.append((start_offset, length))
+    if start_offset + length == len(all_doc_tokens):
+      break
+    start_offset += min(length, 128)
+
+  outputs = []
+  for di, (doc_start, doc_length) in enumerate(doc_spans):
+    tokens = []
+    token_to_orig_map = {}
+    token_is_max_context = {}
+    segment_ids = []
+    tokens.append("[CLS]")
+    segment_ids.append(0)
+    for token in query_tokens:
+      tokens.append(token)
+      segment_ids.append(0)
+    tokens.append("[SEP]")
+    segment_ids.append(0)
+
+    for i in range(doc_length):
+      split_token_index = doc_start + i
+      token_to_orig_map[len(tokens)] = tok_to_orig_index[split_token_index]
+      token_is_max_context[len(tokens)] = _check_is_max_context(doc_spans, di, split_token_index)
+      tokens.append(all_doc_tokens[split_token_index])
+      segment_ids.append(1)
+    tokens.append("[SEP]")
+    segment_ids.append(1)
+
+    input_ids = tokenizer.convert_tokens_to_ids(tokens)
+    input_mask = [1] * len(input_ids)
+
+    while len(input_ids) < 384:
+      input_ids.append(0)
+      input_mask.append(0)
+      segment_ids.append(0)
+
+    assert len(input_ids) == 384
+    assert len(input_mask) == 384
+    assert len(segment_ids) == 384
+
+    outputs.append({
+      "input_ids": np.expand_dims(np.array(input_ids), 0).astype(np.float32),
+      "input_mask": np.expand_dims(np.array(input_mask), 0).astype(np.float32),
+      "segment_ids": np.expand_dims(np.array(segment_ids), 0).astype(np.float32),
+      "token_to_orig_map": token_to_orig_map,
+      "token_is_max_context": token_is_max_context,
+      "tokens": tokens,
+    })
+
+  return outputs
+
+def iterate(tokenizer, start=0):
+  examples = init_dataset()
+  print(f"there are {len(examples)} pairs in the dataset")
+
+  for i in range(start, len(examples)):
+    example = examples[i]
+    features = convert_example_to_features(example, tokenizer)
+    # we need to yield all features here as the f1 score is the maximum over all features
+    yield features, example
+
+if __name__ == "__main__":
+  tokenizer = BertTokenizer(str(Path(__file__).parents[2] / "weights" / "bert_vocab.txt"))
+
+  X, Y = next(iterate(tokenizer))
+  print(" ".join(X[0]["tokens"]))
+  print(X[0]["input_ids"].shape, Y)
--- a/tinygrad_repo/extra/datasets/wikipedia.py
+++ b/tinygrad_repo/extra/datasets/wikipedia.py
@@ -0,0 +1,398 @@
+# Preprocessing of downloaded text from Wikipedia for MLPerf BERT training
+# This is a modified version of the original script:
+# https://github.com/mlcommons/training/blob/master/language_model/tensorflow/bert/cleanup_scripts/create_pretraining_data.py
+# ENV VARS:
+# MAX_SEQ_LENGTH          - Maximum sequence length
+# MAX_PREDICTIONS_PER_SEQ - Maximum number of masked LM predictions per sequence
+# RANDOM_SEED             - Random seed
+# DUPE_FACTOR             - Number of times to duplicate the input data with different masks
+# MASKED_LM_PROB          - Probability of masking a token
+# SHORT_SEQ_PROB          - Probability of picking a sequence shorter than MAX_SEQ_LENGTH
+
+import os, sys, pickle, random, unicodedata
+from pathlib import Path
+import numpy as np
+from tqdm import tqdm
+from tqdm.contrib.concurrent import process_map
+
+from tinygrad.helpers import diskcache, getenv
+
+BASEDIR = getenv('BASEDIR', Path(__file__).parent / "wiki")
+
+################### Tokenization #####################
+
+def _is_whitespace(char:str) -> bool:
+  if char == " " or char == "\t" or char == "\n" or char == "\r":
+    return True
+  return unicodedata.category(char) == "Zs"
+
+def _is_control(char:str) -> bool:
+  if char == "\t" or char == "\n" or char == "\r":
+    return False
+  return unicodedata.category(char).startswith("C")
+
+def _is_punctuation(char:str) -> bool:
+  # range(33, 48) -> ! " # $ % & ' ( ) * + , - . /
+  # range(58, 65) -> : ; < = > ? @
+  # range(91, 97) -> [ \ ] ^ _
+  # range(123, 127) -> { | } ~
+  if (cp := ord(char)) in range(33, 48) or cp in range(58, 65) or cp in range(91, 97) or cp in range(123, 127):
+    return True
+  return unicodedata.category(char).startswith("P")
+
+def _is_chinese_char(cp:int) -> bool:
+  if ((cp >= 0x4E00 and cp <= 0x9FFF) or
+      (cp >= 0x3400 and cp <= 0x4DBF) or
+      (cp >= 0x20000 and cp <= 0x2A6DF) or
+      (cp >= 0x2A700 and cp <= 0x2B73F) or
+      (cp >= 0x2B740 and cp <= 0x2B81F) or
+      (cp >= 0x2B820 and cp <= 0x2CEAF) or
+      (cp >= 0xF900 and cp <= 0xFAFF) or
+      (cp >= 0x2F800 and cp <= 0x2FA1F)):
+    return True
+  return False
+
+def _run_split_on_punc(text:str) -> list[str]:
+  if text in ("[UNK]", "[SEP]", "[PAD]", "[CLS]", "[MASK]"):
+    return [text]
+  start_new_word = True
+  output = []
+  for i in range(len(text)):
+    if _is_punctuation(char := text[i]):
+      output.append([char])
+      start_new_word = True
+    else:
+      if start_new_word:
+        output.append([])
+      start_new_word = False
+      output[-1].append(char)
+  return ["".join(x) for x in output]
+
+def _run_strip_accents(text:str) -> str:
+  output = []
+  for char in unicodedata.normalize("NFD", text):
+    if unicodedata.category(char) != "Mn":
+      output.append(char)
+  return "".join(output)
+
+def _clean_text(text:str) -> str:
+  output = []
+  for char in text:
+    if not ((cp := ord(char)) == 0 or cp == 0xfffd or _is_control(char)):
+      output.append(" " if _is_whitespace(char) else char)
+  return "".join(output)
+
+def _tokenize_chinese_chars(text:str) -> str:
+  output = []
+  for char in text:
+    cp = ord(char)
+    if _is_chinese_char(cp):
+      output.append(" ")
+      output.append(char)
+      output.append(" ")
+    else:
+      output.append(char)
+  return "".join(output)
+
+def whitespace_tokenize(text):
+  if not (text := text.strip()): return []
+  return text.split()
+
+def _wordpiece_tokenize(text:str, vocab:dict[str, int]) -> list[str]:
+  text = text.decode("utf-8", "ignore") if isinstance(text, bytes) else text
+  output_tokens = []
+  for token in text.strip().split():
+    chars = list(token)
+    if len(chars) > 200:
+      output_tokens.append("[UNK]")
+      continue
+
+    is_bad = False
+    start = 0
+    sub_tokens = []
+    while start < len(chars):
+      end = len(chars)
+      cur_substr = None
+      while start < end:
+        substr = "".join(chars[start:end])
+        if start > 0: substr = "##" + substr
+        if substr in vocab:
+          cur_substr = substr
+          break
+        end -= 1
+      if cur_substr is None:
+        is_bad = True
+        break
+      sub_tokens.append(cur_substr)
+      start = end
+
+    if is_bad: output_tokens.append("[UNK]")
+    else: output_tokens.extend(sub_tokens)
+  return output_tokens
+
+class Tokenizer:
+  def __init__(self, vocab_file):
+    self.vocab = {}
+    with open(vocab_file) as f:
+      for line in f:
+        line = line.decode("utf-8", "ignore") if isinstance(line, bytes) else line
+        if (token := line.strip()) and token not in self.vocab: self.vocab[token] = len(self.vocab)
+    self.inv_vocab = {v: k for k, v in self.vocab.items()}
+
+  def tokenize(self, text:str) -> list[str]:
+    # BasicTokenizer
+    split_tokens = []
+    for token in whitespace_tokenize(_tokenize_chinese_chars(_clean_text(text.decode("utf-8", "ignore") if isinstance(text, bytes) else text))):
+      split_tokens.extend(_run_split_on_punc(_run_strip_accents(token.lower())))
+    split_tokens = " ".join(split_tokens).strip().split()
+    # WordpieceTokenizer
+    tokens = []
+    for token in split_tokens:
+      tokens.extend(_wordpiece_tokenize(token, self.vocab))
+    return tokens
+
+  def convert_tokens_to_ids(self, tokens:list[str]) -> list[int]: return [self.vocab[token] for token in tokens]
+  def convert_ids_to_tokens(self, ids:list[int]) -> list[str]: return [self.inv_vocab[id] for id in ids]
+
+##################### Feature transformation #####################
+
+def truncate_seq_pair(tokens_a:list[str], tokens_b:list[str], max_num_tokens:int, rng:random.Random) -> None:
+  while True:
+    total_length = len(tokens_a) + len(tokens_b)
+    if total_length <= max_num_tokens:
+      break
+
+    trunc_tokens = tokens_a if len(tokens_a) > len(tokens_b) else tokens_b
+    assert len(trunc_tokens) >= 1
+
+    if rng.random() < 0.5:
+      del trunc_tokens[0]
+    else:
+      trunc_tokens.pop()
+
+def create_masked_lm_predictions(tokens:list[str], tokenizer:Tokenizer, rng:random.Random, vocab_words:list[str]) -> tuple[list[str], list[int], list[str]]:
+  cand_indices = []
+  for i, token in enumerate(tokens):
+    if token == "[CLS]" or token == "[SEP]":
+      continue
+    cand_indices.append(i)
+
+  rng.shuffle(cand_indices)
+  output_tokens = list(tokens)
+  num_to_predict = min(getenv('MAX_PREDICTIONS_PER_SEQ', 76), max(1, int(round(len(tokens) * getenv("MASKED_LM_PROB", 0.15)))))
+
+  masked_lms = []
+  covered_indices = set()
+  for index in cand_indices:
+    if len(masked_lms) >= num_to_predict:
+      break
+    if index in covered_indices:
+      continue
+    covered_indices.add(index)
+
+    masked_token = None
+    if rng.random() < 0.8:
+      masked_token = "[MASK]"
+    else:
+      if rng.random() < 0.5:
+        masked_token = tokens[index]
+      else:
+        masked_token = vocab_words[rng.randint(0, len(tokenizer.vocab) - 1)]
+
+    output_tokens[index] = masked_token
+    masked_lms.append((index, tokens[index]))
+  masked_lms = sorted(masked_lms, key=lambda x: x[0])
+
+  masked_lm_positions = []
+  masked_lm_labels = []
+  for p in masked_lms:
+    masked_lm_positions.append(p[0])
+    masked_lm_labels.append(p[1])
+
+  return output_tokens, masked_lm_positions, masked_lm_labels
+
+def create_instances_from_document(rng:random.Random, tokenizer:Tokenizer, doc:list[str], di:int, documents:list[list[str]]) -> list[dict]:
+  max_num_tokens = getenv('MAX_SEQ_LENGTH', 512) - 3 # [CLS] + 2 * [SEP]
+
+  target_seq_length = max_num_tokens
+  if rng.random() < getenv("SHORT_SEQ_PROB", 0.1):
+    target_seq_length = rng.randint(2, max_num_tokens)
+
+  instances = []
+  current_chunk = []
+  current_length = 0
+  i = 0
+  while i < len(doc):
+    segment = doc[i]
+    current_chunk.append(segment)
+    current_length += len(segment)
+    if i == len(doc) - 1 or current_length >= target_seq_length:
+      if current_chunk:
+        a_end = 1
+        if len(current_chunk) >= 2:
+          a_end = rng.randint(1, len(current_chunk) - 1)
+
+        tokens_a = []
+        for j in range(a_end):
+          tokens_a.extend(current_chunk[j])
+
+        tokens_b = []
+        is_random_next = False
+        if len(current_chunk) == 1 or rng.random() < 0.5:
+          is_random_next = True
+          target_b_length = target_seq_length - len(tokens_a)
+
+          for _ in range(10):
+            random_document_index = rng.randint(0, len(documents) - 1)
+            if random_document_index != di:
+              break
+
+          random_document = documents[random_document_index]
+          random_start = rng.randint(0, len(random_document) - 1)
+          for j in range(random_start, len(random_document)):
+            tokens_b.extend(random_document[j])
+            if len(tokens_b) >= target_b_length:
+              break
+
+          num_unused_segments = len(current_chunk) - a_end
+          i -= num_unused_segments
+        else:
+          is_random_next = False
+          for j in range(a_end, len(current_chunk)):
+            tokens_b.extend(current_chunk[j])
+        truncate_seq_pair(tokens_a, tokens_b, max_num_tokens, rng)
+
+        assert len(tokens_a) >= 1
+        assert len(tokens_b) >= 1
+
+        tokens = []
+        segment_ids = []
+        tokens.append("[CLS]")
+        segment_ids.append(0)
+        for token in tokens_a:
+          tokens.append(token)
+          segment_ids.append(0)
+        tokens.append("[SEP]")
+        segment_ids.append(0)
+        for token in tokens_b:
+          tokens.append(token)
+          segment_ids.append(1)
+        tokens.append("[SEP]")
+        segment_ids.append(1)
+
+        tokens, masked_lm_positions, masked_lm_labels = create_masked_lm_predictions(tokens, tokenizer, rng, list(tokenizer.vocab.keys()))
+        instances.append({
+          "tokens": tokens,
+          "segment_ids": segment_ids,
+          "masked_lm_positions": masked_lm_positions,
+          "masked_lm_labels": masked_lm_labels,
+          "is_random_next": is_random_next
+        })
+      current_chunk = []
+      current_length = 0
+    i += 1
+  return instances
+
+def get_documents(rng:random.Random, tokenizer:Tokenizer, fn:str) -> list[list[str]]:
+  documents = [[]]
+  with open(BASEDIR / fn) as f:
+    for line in f.readlines():
+      if not (line := line.decode("utf-8", "ignore") if isinstance(line, bytes) else line): break
+      if not (line := line.strip()): documents.append([])
+      if (tokens := tokenizer.tokenize(line)): documents[-1].append(tokens)
+  documents = [x for x in documents if x]
+  rng.shuffle(documents)
+  return documents
+
+def get_instances(rng:random.Random, tokenizer:Tokenizer, documents:list[list[str]]) -> list[dict]:
+  instances = []
+  for _ in range(getenv('DUPE_FACTOR', 10)):
+    for di, doc in enumerate(documents):
+      instances.extend(create_instances_from_document(rng, tokenizer, doc, di, documents))
+  rng.shuffle(instances)
+  return instances
+
+def instance_to_features(instance:dict, tokenizer:Tokenizer) -> dict:
+  input_ids = tokenizer.convert_tokens_to_ids(instance["tokens"])
+  input_mask = [1] * len(input_ids)
+  segment_ids = instance["segment_ids"]
+
+  max_seq_length = getenv('MAX_SEQ_LENGTH', 512)
+
+  assert len(input_ids) <= max_seq_length
+  while len(input_ids) < max_seq_length:
+    input_ids.append(0)
+    input_mask.append(0)
+    segment_ids.append(0)
+  assert len(input_ids) == max_seq_length
+  assert len(input_mask) == max_seq_length
+  assert len(segment_ids) == max_seq_length
+
+  masked_lm_positions = instance["masked_lm_positions"]
+  masked_lm_ids = tokenizer.convert_tokens_to_ids(instance["masked_lm_labels"])
+  masked_lm_weights = [1.0] * len(masked_lm_ids)
+
+  while len(masked_lm_positions) < getenv("MAX_PREDICTIONS_PER_SEQ", 76):
+    masked_lm_positions.append(0)
+    masked_lm_ids.append(0)
+    masked_lm_weights.append(0.0)
+
+  next_sentence_label = 1 if instance["is_random_next"] else 0
+
+  return {
+    "input_ids": np.expand_dims(np.array(input_ids, dtype=np.int32), 0),
+    "input_mask": np.expand_dims(np.array(input_mask, dtype=np.int32), 0),
+    "segment_ids": np.expand_dims(np.array(segment_ids, dtype=np.int32), 0),
+    "masked_lm_positions": np.expand_dims(np.array(masked_lm_positions, dtype=np.int32), 0),
+    "masked_lm_ids": np.expand_dims(np.array(masked_lm_ids, dtype=np.int32), 0),
+    "masked_lm_weights": np.expand_dims(np.array(masked_lm_weights, dtype=np.float32), 0),
+    "next_sentence_labels": np.expand_dims(np.array([next_sentence_label], dtype=np.int32), 0),
+  }
+
+def process_part(part:int):
+  tokenizer = Tokenizer(getenv("BASEDIR", Path(__file__).parent / "wiki") / "vocab.txt")
+  os.makedirs(BASEDIR / "train", exist_ok=True)
+
+  if os.path.exists(BASEDIR / f"train/{str(part)}.pkl"): return
+  features = get_features_from_part(tokenizer, val=False, part=part)
+  with open(BASEDIR / f"train/{str(part)}.pkl", "wb") as f:
+    pickle.dump(features, f)
+
+def get_features_from_part(tokenizer:Tokenizer, val:bool=False, part:int=0) -> list[dict]: # Convert raw text to masked NSP samples
+  rng = random.Random(getenv('RANDOM_SEED', 12345))
+
+  if val:
+    tqdm.write("Getting samples from dataset")
+    documents = get_documents(rng, tokenizer, "results4/eval.txt")
+    instances = get_instances(rng, tokenizer, documents)
+
+    tqdm.write(f"There are {len(instances)} samples in the dataset")
+    tqdm.write(f"Picking 10000 samples")
+
+    pick_ratio = len(instances) / 10000
+    return [instance_to_features(instances[int(inst*pick_ratio)], tokenizer) for inst in range(10000)]
+  else:
+    documents = get_documents(rng, tokenizer, f"results4/part-{part:05d}-of-00500")
+    instances = get_instances(rng, tokenizer, documents)
+    return [instance_to_features(instance, tokenizer) for instance in instances]
+
+##################### Load files #####################
+
+@diskcache
+def get_wiki_train_files(): return sorted(list((BASEDIR / "train/").glob("*.pkl")))
+
+if __name__ == "__main__":
+  tokenizer = Tokenizer(getenv("BASEDIR", Path(__file__).parent / "wiki") / "vocab.txt")
+
+  assert len(sys.argv) > 1, "Usage: python wikipedia.py pre-eval|pre-train [part]|all"
+
+  if sys.argv[1] == "pre-eval": # Generate 10000 eval samples
+    with open(BASEDIR / "eval.pkl", "wb") as f:
+      pickle.dump(get_features_from_part(tokenizer, val=True), f)
+  elif sys.argv[1] == "pre-train":
+    if sys.argv[2] == "all": # Use all 500 parts for training generation
+      process_map(process_part, [part for part in range(500)], max_workers=getenv('NUM_WORKERS', min(os.cpu_count(), 32)), chunksize=1)
+    else: # Use a specific part for training generation
+      part = sys.argv[2]
+      print(f"Processing part {part}...")
+      process_part(int(part))
--- a/tinygrad_repo/extra/datasets/wikipedia_download.py
+++ b/tinygrad_repo/extra/datasets/wikipedia_download.py
@@ -0,0 +1,54 @@
+# pip install gdown
+# Downloads the 2020 wikipedia dataset used for MLPerf BERT training
+import os, hashlib
+from pathlib import Path
+import tarfile
+import gdown
+from tqdm import tqdm
+from tinygrad.helpers import getenv
+
+def gdrive_download(url:str, path:str):
+  if not os.path.exists(path): gdown.download(url, path)
+
+def wikipedia_uncompress_and_extract(file:str, path:str, small:bool=False):
+  if not os.path.exists(os.path.join(path, "results4")):
+    print("Uncompressing and extracting file...")
+    with tarfile.open(file, 'r:gz') as tar:
+      tar.extractall(path=path)
+      os.remove(file)
+      if small:
+        for member in tar.getmembers(): tar.extract(path=path, member=member)
+      else:
+        for member in tqdm(iterable=tar.getmembers(), total=len(tar.getmembers())): tar.extract(path=path, member=member)
+
+def verify_checksum(folder_path:str, checksum_path:str):
+  print("Verifying checksums...")
+  with open(checksum_path, 'r') as f:
+    for line in f:
+      expected_checksum, folder_name = line.split()
+      file_path = os.path.join(folder_path, folder_name[2:]) # remove './' from the start of the folder name
+      hasher = hashlib.md5()
+      with open(file_path, 'rb') as f:
+        for buf in iter(lambda: f.read(4096), b''): hasher.update(buf)
+      if hasher.hexdigest() != expected_checksum:
+        raise ValueError(f"Checksum does not match for file: {file_path}")
+  print("All checksums match.")
+
+def download_wikipedia(path:str):
+  # Links from: https://github.com/mlcommons/training/blob/master/language_model/tensorflow/bert/dataset.md
+  os.makedirs(path, exist_ok=True)
+  gdrive_download("https://drive.google.com/uc?id=1fbGClQMi2CoMv7fwrwTC5YYPooQBdcFW", os.path.join(path, "bert_config.json"))
+  gdrive_download("https://drive.google.com/uc?id=1USK108J6hMM_d27xCHi738qBL8_BT1u1", os.path.join(path, "vocab.txt"))
+  gdrive_download("https://drive.google.com/uc?id=1chiTBljF0Eh1U5pKs6ureVHgSbtU8OG_", os.path.join(path, "model.ckpt-28252.data-00000-of-00001"))
+  gdrive_download("https://drive.google.com/uc?id=1Q47V3K3jFRkbJ2zGCrKkKk-n0fvMZsa0", os.path.join(path, "model.ckpt-28252.index"))
+  gdrive_download("https://drive.google.com/uc?id=1vAcVmXSLsLeQ1q7gvHnQUSth5W_f_pwv", os.path.join(path, "model.ckpt-28252.meta"))
+  with open(os.path.join(path, "checkpoint"), "w") as f: f.write('model_checkpoint_path: "model.ckpt-28252"\nall_model_checkpoint_paths: "model.ckpt-28252"')
+  if getenv("WIKI_TRAIN", 0):
+    gdrive_download("https://drive.google.com/uc?id=1tmMgLwoBvbEJEHXh77sqrXYw5RpqT8R_", os.path.join(path, "bert_reference_results_text_md5.txt"))
+    gdrive_download("https://drive.google.com/uc?id=14xV2OUGSQDG_yDBrmbSdcDC-QGeqpfs_", os.path.join(path, "results_text.tar.gz"))
+    wikipedia_uncompress_and_extract(os.path.join(path, "results_text.tar.gz"), path)
+    if getenv("VERIFY_CHECKSUM", 0):
+      verify_checksum(os.path.join(path, "results4"), os.path.join(path, "bert_reference_results_text_md5.txt"))
+
+if __name__ == "__main__":
+  download_wikipedia(getenv("BASEDIR", os.path.join(Path(__file__).parent / "wiki")))