Add llms subdir + update README (#145)

* add llms subdir + update README

* nits

* use same pre-commit as mlx

* update readmes a bit

* format

transformer_lm/datasets.py

@@ -2,10 +2,11 @@
 
 import io
 import itertools
-import numpy as np
 import os
-from urllib import request
 import zipfile
+from urllib import request
+
+import numpy as np
 
 
 def load_dataset(dataname):

transformer_lm/jax_main.py

@@ -1,14 +1,14 @@
 # Copyright © 2023 Apple Inc.
 
 import functools
-import jax
-import jax.numpy as jnp
 import math
-import numpy as np
 import time
 from collections import namedtuple
 
 import datasets
+import jax
+import jax.numpy as jnp
+import numpy as np
 from tree_utils import tree_flatten
 
 """

transformer_lm/main.py

@@ -3,15 +3,13 @@
 import math
 import time
 
-import numpy as np
-
+import datasets
 import mlx.core as mx
 import mlx.nn as nn
 import mlx.optimizers as optim
+import numpy as np
 from mlx.utils import tree_flatten
 
-import datasets
-
 
 class TransformerLM(nn.Module):
     def __init__(self, vocab_size: int, num_layers: int, dims: int, num_heads: int):

transformer_lm/tf_main.py

@@ -3,11 +3,10 @@
 import math
 import time
 
+import datasets
 import numpy as np
 import tensorflow as tf
 
-import datasets
-
 
 def to_samples(context_size, dataset):
     tokens = dataset.size

transformer_lm/torch_main.py

@@ -3,11 +3,10 @@
 import math
 import time
 
+import datasets
 import numpy as np
 import torch
 
-import datasets
-
 
 def to_samples(context_size, dataset):
     tokens = dataset.size