bert-base-ner-mlx / README.md
masahiroid's picture
Update links: masahirocom -> masahiroid; install via PyPI (pip install model-audit-lite)
9739570 verified
|
Raw History Blame Contribute Delete
6.53 kB
---
license: mit
language:
- en
- ja
tags:
- mlx
- token-classification
- named-entity-recognition
- bert
base_model: dslim/bert-base-NER
pipeline_tag: token-classification
---
# bert-base-ner-mlx
[English](#model-summary) | [日本語](#モデルの概要)
## Model Summary
This is an unofficial [MLX](https://github.com/ml-explore/mlx) conversion of
[dslim/bert-base-NER](https://huggingface.co/dslim/bert-base-NER) (a standard
BERT-base with a named-entity-recognition head on top, trained on CoNLL-2003:
PER/ORG/LOC/MISC). All credit for the original model goes to its author.
### This cannot be loaded with mlx-embeddings
`mlx-embeddings` targets pooled embedding/reranker outputs and doesn't
support a per-token classification head. This model was **reimplemented from
scratch for MLX** and requires the bundled `bert_ner_mlx.py`. The
architecture itself is simple: a standard post-norm BERT encoder plus a
linear classification head.
### Usage
```python
import mlx.core as mx
from mlx.utils import tree_unflatten
from transformers import AutoTokenizer
from bert_ner_mlx import BertNerMLX
tokenizer = AutoTokenizer.from_pretrained("dslim/bert-base-NER")
model = BertNerMLX()
weights = mx.load("model.safetensors")
model.update(tree_unflatten(list(weights.items())))
mx.eval(model.parameters())
text = "My name is Wolfgang and I live in Berlin, working at Hugging Face."
inputs = tokenizer(text, return_tensors="np")
input_ids = mx.array(inputs["input_ids"])
attention_mask = mx.array(inputs["attention_mask"])
logits = model(input_ids, attention_mask=attention_mask) # (1, seq_len, 9)
labels = logits.argmax(-1)[0].tolist()
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
for tok, label_id in zip(tokens, labels):
print(tok, model.config.id2label[label_id] if hasattr(model, "config") else label_id)
```
Labels are the same as the original model: `{0: "O", 1: "B-MISC", 2: "I-MISC",
3: "B-PER", 4: "I-PER", 5: "B-ORG", 6: "I-ORG", 7: "B-LOC", 8: "I-LOC"}`.
### Accuracy
Compared against the PyTorch fp32 reference on a real sentence containing
person/location/organization entities:
| Precision | Logits cosine sim. | Label agreement |
|---|---|---|
| MLX fp32 | 1.0 | 100% |
| MLX fp16 (this release) | 0.99999994 | 100% |
During fp16 conversion, using an attention-mask value like `-1e9` (which
doesn't fit in fp16) caused `0 * (-inf) = NaN` at unmasked (`mask=1`)
positions. Fixed by switching to `-1e4`, which fits fp16's range (full
details are in the conversion-toolkit repo, not `SECURITY.md`).
### Specs
| Item | Value |
|---|---|
| Base model | `dslim/bert-base-NER` (BERT-base, 108M params) |
| Precision | float16 |
| Framework | MLX (from-scratch `bert_ner_mlx.py`) |
### Notes
- This is a community conversion, not an official release from the original author.
- Security audit uses [model-audit-lite](https://github.com/masahiroid/model-audit-lite)
(see `SECURITY.md` for details).
### Security
Audited against its upstream with [model-audit-lite](https://github.com/masahiroid/model-audit-lite): weight format, bundled code, and a machine-readable lineage (ML-BOM). Details, checksums and how to reproduce: [SECURITY.md](SECURITY.md).
---
## モデルの概要
[dslim/bert-base-NER](https://huggingface.co/dslim/bert-base-NER)(標準的なBERT-baseに
固有表現抽出ヘッドを乗せたモデル。CoNLL-2003、PER/ORG/LOC/MISCの4種)の
[MLX](https://github.com/ml-explore/mlx)版です。元モデルの著作権はその作者に帰属します。
### mlx-embeddingsでは読み込めません
`mlx-embeddings`はembedding/reranker用のpooling出力をターゲットにしており、トークン単位の
分類ヘッド(各トークンに対して独立にラベルを出す構成)には対応していないため、**MLXでの
実装をゼロから書き起こして**変換しています。同梱の`bert_ner_mlx.py`が必要です。構造自体は
標準的なpost-norm BERTエンコーダー + 線形分類ヘッドとシンプルです。
### 使い方
```python
import mlx.core as mx
from mlx.utils import tree_unflatten
from transformers import AutoTokenizer
from bert_ner_mlx import BertNerMLX
tokenizer = AutoTokenizer.from_pretrained("dslim/bert-base-NER")
model = BertNerMLX()
weights = mx.load("model.safetensors")
model.update(tree_unflatten(list(weights.items())))
mx.eval(model.parameters())
text = "My name is Wolfgang and I live in Berlin, working at Hugging Face."
inputs = tokenizer(text, return_tensors="np")
input_ids = mx.array(inputs["input_ids"])
attention_mask = mx.array(inputs["attention_mask"])
logits = model(input_ids, attention_mask=attention_mask) # (1, seq_len, 9)
labels = logits.argmax(-1)[0].tolist()
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
for tok, label_id in zip(tokens, labels):
print(tok, model.config.id2label[label_id] if hasattr(model, "config") else label_id)
```
ラベルは元モデルと同じ: `{0: "O", 1: "B-MISC", 2: "I-MISC", 3: "B-PER", 4: "I-PER",
5: "B-ORG", 6: "I-ORG", 7: "B-LOC", 8: "I-LOC"}`。
### 精度検証
PyTorch fp32リファレンスと、実際の文章1件(人名・地名・組織名を含む)で比較:
| 精度 | Logitsコサイン類似度 | ラベル一致率 |
|---|---|---|
| MLX fp32 | 1.0 | 100% |
| MLX fp16(本リリース) | 0.99999994 | 100% |
fp16変換の際、attention maskのマスク値に`-1e9`のようなfp16で表現できない大きな負数を使うと、
マスクされていない(`mask=1`)位置で`0 * (-inf)`が発生してNaNになる罠があった。
fp16の範囲内に収まる`-1e4`に変更して解消している(詳細は`SECURITY.md`ではなく、変換ノウハウ
リポジトリに記載)。
### Specs
| Item | Value |
|---|---|
| ベースモデル | `dslim/bert-base-NER`(BERT-base、108M params) |
| 精度 | float16 |
| フレームワーク | MLX(ゼロから実装した`bert_ner_mlx.py`) |
### 備考
- 本変換は非公式のコミュニティ版です。
- セキュリティー監査には[model-audit-lite](https://github.com/masahiroid/model-audit-lite)を
使用しています(詳細は`SECURITY.md`)。
### セキュリティー
[model-audit-lite](https://github.com/masahiroid/model-audit-lite) で変換元と突き合わせて監査済みです(重みの形式、同梱コード、機械可読な系譜=ML-BOM)。詳細・チェックサム・再現方法は [SECURITY.md](SECURITY.md) をご覧ください。