Enable more BERT models (#580)

* Update convert.py

* Update model.py

* Update test.py

* Update model.py

* Update convert.py

* Add files via upload

* Update convert.py

* format

* nit

* nit

---------

Co-authored-by: Awni Hannun <awni@apple.com>

bert/test.py

@@ -1,3 +1,4 @@
+import argparse
 from typing import List
 
 import model
@@ -16,19 +17,41 @@ def run_torch(bert_model: str, batch: List[str]):
 
 
 if __name__ == "__main__":
-    bert_model = "bert-base-uncased"
-    mlx_model = "weights/bert-base-uncased.npz"
-    batch = [
-        "This is an example of BERT working in MLX.",
-        "A second string",
-        "This is another string.",
-    ]
-    torch_output, torch_pooled = run_torch(bert_model, batch)
-    mlx_output, mlx_pooled = model.run(bert_model, mlx_model, batch)
-    assert np.allclose(
-        torch_output, mlx_output, rtol=1e-4, atol=1e-5
-    ), "Model output is different"
-    assert np.allclose(
-        torch_pooled, mlx_pooled, rtol=1e-4, atol=1e-5
-    ), "Model pooled output is different"
-    print("Tests pass :)")
+    parser = argparse.ArgumentParser(
+        description="Run a BERT-like model for a batch of text."
+    )
+    parser.add_argument(
+        "--bert-model",
+        type=str,
+        default="bert-base-uncased",
+        help="The model identifier for a BERT-like model from Hugging Face Transformers.",
+    )
+    parser.add_argument(
+        "--mlx-model",
+        type=str,
+        default="weights/bert-base-uncased.npz",
+        help="The path of the stored MLX BERT weights (npz file).",
+    )
+    parser.add_argument(
+        "--text",
+        nargs="+",
+        default=["This is an example of BERT working in MLX."],
+        help="A batch of texts to process. Multiple texts should be separated by spaces.",
+    )
+
+    args = parser.parse_args()
+
+    torch_output, torch_pooled = run_torch(args.bert_model, args.text)
+
+    mlx_output, mlx_pooled = model.run(args.bert_model, args.mlx_model, args.text)
+
+    if torch_pooled is not None and mlx_pooled is not None:
+        assert np.allclose(
+            torch_output, mlx_output, rtol=1e-4, atol=1e-5
+        ), "Model output is different"
+        assert np.allclose(
+            torch_pooled, mlx_pooled, rtol=1e-4, atol=1e-5
+        ), "Model pooled output is different"
+        print("Tests pass :)")
+    else:
+        print("Pooled outputs were not compared due to one or both being None.")