Upload 3 files

convert.py

@@ -0,0 +1,66 @@
+import argparse
+import os
+from pathlib import Path
+
+import numpy
+from transformers import AutoModel, AutoConfig
+
+
+def replace_key(key: str) -> str:
+    key = key.replace(".layer.", ".layers.")
+    key = key.replace(".self.key.", ".key_proj.")
+    key = key.replace(".self.query.", ".query_proj.")
+    key = key.replace(".self.value.", ".value_proj.")
+    key = key.replace(".attention.output.dense.", ".attention.out_proj.")
+    key = key.replace(".attention.output.LayerNorm.", ".ln1.")
+    key = key.replace(".output.LayerNorm.", ".ln2.")
+    key = key.replace(".intermediate.dense.", ".linear1.")
+    key = key.replace(".output.dense.", ".linear2.")
+    key = key.replace(".LayerNorm.", ".norm.")
+    key = key.replace("pooler.dense.", "pooler.")
+    return key
+
+
+def convert(bert_model: str, mlx_model: str) -> None:
+    # Load model and its configuration
+    model = AutoModel.from_pretrained(bert_model)
+    config = AutoConfig.from_pretrained(bert_model)
+    
+    # Create output directory if it doesn't exist
+    output_dir = os.path.dirname(mlx_model)
+    if output_dir and not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+    
+    # Save config as well
+    config_path = os.path.join(output_dir, "config.json")
+    with open(config_path, "w") as f:
+        f.write(config.to_json_string())
+        
+    print(f"Saved model config to {config_path}")
+    
+    # Save the tensors
+    tensors = {
+        replace_key(key): tensor.numpy() for key, tensor in model.state_dict().items()
+    }
+    numpy.savez(mlx_model, **tensors)
+    print(f"Saved model weights to {mlx_model}")
+    print(f"Model vocab size: {config.vocab_size}, hidden size: {config.hidden_size}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert BERT weights to MLX.")
+    parser.add_argument(
+        "--bert-model",
+        type=str,
+        default="bert-base-uncased",
+        help="The huggingface name of the BERT model to save. Any BERT-like model can be specified.",
+    )
+    parser.add_argument(
+        "--mlx-model",
+        type=str,
+        default="weights/bert-base-uncased.npz",
+        help="The output path for the MLX BERT weights.",
+    )
+    args = parser.parse_args()
+
+    convert(args.bert_model, args.mlx_model)

model.py

@@ -0,0 +1,203 @@
+import argparse
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy
+import numpy as np
+from mlx.utils import tree_unflatten
+from transformers import AutoConfig, AutoTokenizer, PreTrainedTokenizerBase
+
+
+class TransformerEncoderLayer(nn.Module):
+    """
+    A transformer encoder layer with (the original BERT) post-normalization.
+    """
+
+    def __init__(
+        self,
+        dims: int,
+        num_heads: int,
+        mlp_dims: Optional[int] = None,
+        layer_norm_eps: float = 1e-12,
+    ):
+        super().__init__()
+        mlp_dims = mlp_dims or dims * 4
+        self.attention = nn.MultiHeadAttention(dims, num_heads, bias=True)
+        self.ln1 = nn.LayerNorm(dims, eps=layer_norm_eps)
+        self.ln2 = nn.LayerNorm(dims, eps=layer_norm_eps)
+        self.linear1 = nn.Linear(dims, mlp_dims)
+        self.linear2 = nn.Linear(mlp_dims, dims)
+        self.gelu = nn.GELU()
+
+    def __call__(self, x, mask):
+        attention_out = self.attention(x, x, x, mask)
+        add_and_norm = self.ln1(x + attention_out)
+
+        ff = self.linear1(add_and_norm)
+        ff_gelu = self.gelu(ff)
+        ff_out = self.linear2(ff_gelu)
+        x = self.ln2(ff_out + add_and_norm)
+
+        return x
+
+
+class TransformerEncoder(nn.Module):
+    def __init__(
+        self, num_layers: int, dims: int, num_heads: int, mlp_dims: Optional[int] = None
+    ):
+        super().__init__()
+        self.layers = [
+            TransformerEncoderLayer(dims, num_heads, mlp_dims)
+            for i in range(num_layers)
+        ]
+
+    def __call__(self, x, mask):
+        for layer in self.layers:
+            x = layer(x, mask)
+
+        return x
+
+
+class BertEmbeddings(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.token_type_embeddings = nn.Embedding(
+            config.type_vocab_size, config.hidden_size
+        )
+        self.position_embeddings = nn.Embedding(
+            config.max_position_embeddings, config.hidden_size
+        )
+        self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def __call__(
+        self, input_ids: mx.array, token_type_ids: mx.array = None
+    ) -> mx.array:
+        words = self.word_embeddings(input_ids)
+        position = self.position_embeddings(
+            mx.broadcast_to(mx.arange(input_ids.shape[1]), input_ids.shape)
+        )
+
+        if token_type_ids is None:
+            # If token_type_ids is not provided, default to zeros
+            token_type_ids = mx.zeros_like(input_ids)
+
+        token_types = self.token_type_embeddings(token_type_ids)
+
+        embeddings = position + words + token_types
+        return self.norm(embeddings)
+
+
+class Bert(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.embeddings = BertEmbeddings(config)
+        self.encoder = TransformerEncoder(
+            num_layers=config.num_hidden_layers,
+            dims=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            mlp_dims=config.intermediate_size,
+        )
+        self.pooler = nn.Linear(config.hidden_size, config.hidden_size)
+
+    def __call__(
+        self,
+        input_ids: mx.array,
+        token_type_ids: mx.array = None,
+        attention_mask: mx.array = None,
+    ) -> Tuple[mx.array, mx.array]:
+        x = self.embeddings(input_ids, token_type_ids)
+
+        if attention_mask is not None:
+            # convert 0's to -infs, 1's to 0's, and make it broadcastable
+            attention_mask = mx.log(attention_mask)
+            attention_mask = mx.expand_dims(attention_mask, (1, 2))
+
+        y = self.encoder(x, attention_mask)
+        return y, mx.tanh(self.pooler(y[:, 0]))
+
+
+def load_model(
+    bert_model: str, weights_path: str
+) -> Tuple[Bert, PreTrainedTokenizerBase]:
+    if not Path(weights_path).exists():
+        raise ValueError(f"No model weights found in {weights_path}")
+    
+    # First check if there's a local config
+    config_path = Path(weights_path).parent / "config.json"
+    if config_path.exists():
+        with open(config_path, "r") as f:
+            config_dict = json.load(f)
+        config = AutoConfig.for_model(**config_dict)
+        print(f"Loaded local config from {config_path}")
+    else:
+        # If no local config, use the HuggingFace one
+        config = AutoConfig.from_pretrained(bert_model)
+        print(f"Loaded config from HuggingFace for {bert_model}")
+
+    # Create and update the model
+    print(f"Creating model with vocab_size={config.vocab_size}, hidden_size={config.hidden_size}")
+    model = Bert(config)
+    model.load_weights(weights_path)
+
+    tokenizer = AutoTokenizer.from_pretrained(bert_model)
+
+    return model, tokenizer
+
+
+def run(bert_model: str, mlx_model: str, batch: List[str]):
+    import time
+    
+    # Time model loading
+    load_start = time.time()
+    model, tokenizer = load_model(bert_model, mlx_model)
+    load_time = time.time() - load_start
+    print(f"[MLX] Model loaded in {load_time:.2f} seconds")
+    
+    # Time tokenization
+    print(f"[MLX] Tokenizing batch of {len(batch)} sentences")
+    token_start = time.time()
+    tokens = tokenizer(batch, return_tensors="np", padding=True)
+    token_time = time.time() - token_start
+    print(f"[MLX] Tokenization completed in {token_time:.4f} seconds")
+    
+    print(f"[MLX] Tokens shape: input_ids={tokens['input_ids'].shape}")
+    tokens = {key: mx.array(v) for key, v in tokens.items()}
+    
+    # Time inference
+    print(f"[MLX] Running model inference")
+    infer_start = time.time()
+    output, pooled = model(**tokens)
+    mx.eval(output, pooled)  # Force evaluation of lazy arrays
+    infer_time = time.time() - infer_start
+    print(f"[MLX] Inference completed in {infer_time:.4f} seconds")
+    
+    return output, pooled
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Run the BERT model using MLX.")
+    parser.add_argument(
+        "--bert-model",
+        type=str,
+        default="bert-base-uncased",
+        help="The huggingface name of the BERT model to save.",
+    )
+    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",
+        type=str,
+        default="This is an example of BERT working in MLX",
+        help="The text to generate embeddings for.",
+    )
+    args = parser.parse_args()
+    run(args.bert_model, args.mlx_model, args.text)

test.py

@@ -0,0 +1,143 @@
+import argparse
+import time
+from typing import List
+
+import model
+import numpy as np
+import mlx.core as mx
+from transformers import AutoModel, AutoTokenizer
+
+
+def run_torch(bert_model: str, batch: List[str]):
+    print(f"\n[PyTorch] Loading model and tokenizer: {bert_model}")
+    start_time = time.time()
+    tokenizer = AutoTokenizer.from_pretrained(bert_model)
+    torch_model = AutoModel.from_pretrained(bert_model)
+    load_time = time.time() - start_time
+    print(f"[PyTorch] Model loaded in {load_time:.2f} seconds")
+    
+    print(f"[PyTorch] Tokenizing batch of {len(batch)} sentences")
+    torch_tokens = tokenizer(batch, return_tensors="pt", padding=True)
+    
+    print(f"[PyTorch] Running model inference")
+    inference_start = time.time()
+    torch_forward = torch_model(**torch_tokens)
+    inference_time = time.time() - inference_start
+    print(f"[PyTorch] Inference completed in {inference_time:.4f} seconds")
+    
+    torch_output = torch_forward.last_hidden_state.detach().numpy()
+    torch_pooled = torch_forward.pooler_output.detach().numpy()
+    
+    print(f"[PyTorch] Output shape: {torch_output.shape}")
+    print(f"[PyTorch] Pooled output shape: {torch_pooled.shape}")
+    
+    # Print a small sample of the output to verify sensible values
+    print(f"[PyTorch] Sample of output (first token, first 5 values): {torch_output[0, 0, :5]}")
+    print(f"[PyTorch] Sample of pooled output (first 5 values): {torch_pooled[0, :5]}")
+    
+    return torch_output, torch_pooled
+
+
+def run_mlx(bert_model: str, mlx_model: str, batch: List[str]):
+    print(f"\n[MLX] Loading model and tokenizer with weights from: {mlx_model}")
+    start_time = time.time()
+    mlx_output, mlx_pooled = model.run(bert_model, mlx_model, batch)
+    load_and_run_time = time.time() - start_time
+    print(f"[MLX] Model loaded and run in {load_and_run_time:.2f} seconds")
+    
+    # Convert from MLX arrays to numpy for comparison
+    # The correct way to convert MLX arrays to numpy
+    mlx_output_np = np.array(mlx_output)
+    mlx_pooled_np = np.array(mlx_pooled)
+    
+    print(f"[MLX] Output shape: {mlx_output_np.shape}")
+    print(f"[MLX] Pooled output shape: {mlx_pooled_np.shape}")
+    
+    # Print a small sample of the output to verify sensible values
+    print(f"[MLX] Sample of output (first token, first 5 values): {mlx_output_np[0, 0, :5]}")
+    print(f"[MLX] Sample of pooled output (first 5 values): {mlx_pooled_np[0, :5]}")
+    
+    return mlx_output_np, mlx_pooled_np
+
+
+def compare_outputs(torch_output, torch_pooled, mlx_output, mlx_pooled):
+    print("\n[Comparison] Comparing PyTorch and MLX outputs")
+    
+    # Check shapes
+    print(f"[Comparison] Shape match - Output: {torch_output.shape == mlx_output.shape}")
+    print(f"[Comparison] Shape match - Pooled: {torch_pooled.shape == mlx_pooled.shape}")
+    
+    # Calculate differences
+    output_max_diff = np.max(np.abs(torch_output - mlx_output))
+    output_mean_diff = np.mean(np.abs(torch_output - mlx_output))
+    pooled_max_diff = np.max(np.abs(torch_pooled - mlx_pooled))
+    pooled_mean_diff = np.mean(np.abs(torch_pooled - mlx_pooled))
+    
+    print(f"[Comparison] Output - Max absolute difference: {output_max_diff:.6f}")
+    print(f"[Comparison] Output - Mean absolute difference: {output_mean_diff:.6f}")
+    print(f"[Comparison] Pooled - Max absolute difference: {pooled_max_diff:.6f}")
+    print(f"[Comparison] Pooled - Mean absolute difference: {pooled_mean_diff:.6f}")
+    
+    # Detailed comparison of first few values from first sentence
+    print("\n[Comparison] Detailed comparison of first 5 values from first output token:")
+    for i in range(5):
+        torch_val = torch_output[0, 0, i]
+        mlx_val = mlx_output[0, 0, i]
+        diff = abs(torch_val - mlx_val)
+        print(f"Index {i}: PyTorch={torch_val:.6f}, MLX={mlx_val:.6f}, Diff={diff:.6f}")
+    
+    # Check if outputs are close
+    outputs_close = np.allclose(torch_output, mlx_output, rtol=1e-4, atol=1e-4)
+    pooled_close = np.allclose(torch_pooled, mlx_pooled, rtol=1e-4, atol=1e-4)
+    
+    print(f"\n[Comparison] Outputs match within tolerance: {outputs_close}")
+    print(f"[Comparison] Pooled outputs match within tolerance: {pooled_close}")
+    
+    return outputs_close and pooled_close
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Run a BERT-like model for a batch of text and compare PyTorch and MLX outputs."
+    )
+    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.",
+    )
+    parser.add_argument(
+        "--verbose",
+        action="store_true",
+        help="Print detailed information about the model execution.",
+    )
+
+    args = parser.parse_args()
+    
+    print(f"Testing BERT model: {args.bert_model}")
+    print(f"MLX weights: {args.mlx_model}")
+    print(f"Input text: {args.text}")
+    
+    # Run both implementations
+    torch_output, torch_pooled = run_torch(args.bert_model, args.text)
+    mlx_output, mlx_pooled = run_mlx(args.bert_model, args.mlx_model, args.text)
+    
+    # Compare outputs
+    all_match = compare_outputs(torch_output, torch_pooled, mlx_output, mlx_pooled)
+    
+    if all_match:
+        print("\n✅ TEST PASSED: PyTorch and MLX implementations produce equivalent results!")
+    else:
+        print("\n❌ TEST FAILED: PyTorch and MLX implementations produce different results.")