Add/update rampart_mlx.py

rampart_mlx.py

@@ -0,0 +1,155 @@
+"""Rampart BERT-for-token-classification in MLX.
+
+Mirrors HuggingFace `BertForTokenClassification` (MiniLM-L6-H384) so the float
+weights extracted from the ONNX model load 1:1, and so the same module can be
+quantized with mlx and re-exported for mlx-swift.
+"""
+import json
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class BertConfig:
+    def __init__(self, **kw):
+        self.vocab_size = kw["vocab_size"]
+        self.hidden_size = kw["hidden_size"]
+        self.num_hidden_layers = kw["num_hidden_layers"]
+        self.num_attention_heads = kw["num_attention_heads"]
+        self.intermediate_size = kw["intermediate_size"]
+        self.max_position_embeddings = kw["max_position_embeddings"]
+        self.type_vocab_size = kw["type_vocab_size"]
+        self.layer_norm_eps = kw.get("layer_norm_eps", 1e-12)
+        self.num_labels = len(kw["id2label"])
+        self.id2label = {int(k): v for k, v in kw["id2label"].items()}
+
+    @classmethod
+    def from_json(cls, path):
+        with open(path) as f:
+            return cls(**json.load(f))
+
+
+class BertEmbeddings(nn.Module):
+    def __init__(self, c: BertConfig):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(c.vocab_size, c.hidden_size)
+        self.position_embeddings = nn.Embedding(c.max_position_embeddings, c.hidden_size)
+        self.token_type_embeddings = nn.Embedding(c.type_vocab_size, c.hidden_size)
+        self.LayerNorm = nn.LayerNorm(c.hidden_size, eps=c.layer_norm_eps)
+
+    def __call__(self, input_ids, token_type_ids):
+        seq = input_ids.shape[1]
+        pos = mx.arange(seq)[None, :]
+        e = (self.word_embeddings(input_ids)
+             + self.position_embeddings(pos)
+             + self.token_type_embeddings(token_type_ids))
+        return self.LayerNorm(e)
+
+
+class BertSelfAttention(nn.Module):
+    def __init__(self, c: BertConfig):
+        super().__init__()
+        self.num_heads = c.num_attention_heads
+        self.head_dim = c.hidden_size // c.num_attention_heads
+        self.query = nn.Linear(c.hidden_size, c.hidden_size)
+        self.key = nn.Linear(c.hidden_size, c.hidden_size)
+        self.value = nn.Linear(c.hidden_size, c.hidden_size)
+
+    def __call__(self, x, mask):
+        B, L, _ = x.shape
+        H, D = self.num_heads, self.head_dim
+        def shape(t):
+            return t.reshape(B, L, H, D).transpose(0, 2, 1, 3)
+        q, k, v = shape(self.query(x)), shape(self.key(x)), shape(self.value(x))
+        scores = (q @ k.transpose(0, 1, 3, 2)) * (1.0 / (D ** 0.5))
+        scores = scores + mask  # mask: [B,1,1,L] additive
+        probs = mx.softmax(scores, axis=-1)
+        ctx = (probs @ v).transpose(0, 2, 1, 3).reshape(B, L, H * D)
+        return ctx
+
+
+class BertAttention(nn.Module):
+    def __init__(self, c: BertConfig):
+        super().__init__()
+        self.self = BertSelfAttention(c)
+        self.output = BertSelfOutput(c)
+
+    def __call__(self, x, mask):
+        return self.output(self.self(x, mask), x)
+
+
+class BertSelfOutput(nn.Module):
+    def __init__(self, c: BertConfig):
+        super().__init__()
+        self.dense = nn.Linear(c.hidden_size, c.hidden_size)
+        self.LayerNorm = nn.LayerNorm(c.hidden_size, eps=c.layer_norm_eps)
+
+    def __call__(self, x, residual):
+        return self.LayerNorm(self.dense(x) + residual)
+
+
+class BertIntermediate(nn.Module):
+    def __init__(self, c: BertConfig):
+        super().__init__()
+        self.dense = nn.Linear(c.hidden_size, c.intermediate_size)
+
+    def __call__(self, x):
+        return nn.gelu(self.dense(x))
+
+
+class BertOutput(nn.Module):
+    def __init__(self, c: BertConfig):
+        super().__init__()
+        self.dense = nn.Linear(c.intermediate_size, c.hidden_size)
+        self.LayerNorm = nn.LayerNorm(c.hidden_size, eps=c.layer_norm_eps)
+
+    def __call__(self, x, residual):
+        return self.LayerNorm(self.dense(x) + residual)
+
+
+class BertLayer(nn.Module):
+    def __init__(self, c: BertConfig):
+        super().__init__()
+        self.attention = BertAttention(c)
+        self.intermediate = BertIntermediate(c)
+        self.output = BertOutput(c)
+
+    def __call__(self, x, mask):
+        a = self.attention(x, mask)
+        return self.output(self.intermediate(a), a)
+
+
+class BertEncoder(nn.Module):
+    def __init__(self, c: BertConfig):
+        super().__init__()
+        self.layer = [BertLayer(c) for _ in range(c.num_hidden_layers)]
+
+    def __call__(self, x, mask):
+        for lyr in self.layer:
+            x = lyr(x, mask)
+        return x
+
+
+class BertModel(nn.Module):
+    def __init__(self, c: BertConfig):
+        super().__init__()
+        self.embeddings = BertEmbeddings(c)
+        self.encoder = BertEncoder(c)
+
+    def __call__(self, input_ids, token_type_ids, mask):
+        h = self.embeddings(input_ids, token_type_ids)
+        return self.encoder(h, mask)
+
+
+class RampartForTokenClassification(nn.Module):
+    def __init__(self, c: BertConfig):
+        super().__init__()
+        self.config = c
+        self.bert = BertModel(c)
+        self.classifier = nn.Linear(c.hidden_size, c.num_labels)
+
+    def __call__(self, input_ids, attention_mask, token_type_ids=None):
+        if token_type_ids is None:
+            token_type_ids = mx.zeros_like(input_ids)
+        mask = (1.0 - attention_mask.astype(mx.float32))[:, None, None, :] * -1e9
+        h = self.bert(input_ids, token_type_ids, mask)
+        return self.classifier(h)