EntityBERT: Lightweight Named Entity Recognition for Edge AI and IoT

Overview

EntityBERT is a fine-tuned, lightweight transformer model built on boltuix/bert-mini for Named Entity Recognition (NER). With ~4.4M parameters and a size of ~15 MB, it identifies 36 entity types (e.g., PERSON, ORG, GPE, DATE) in English text, optimized for edge AI, IoT devices, and mobile applications. Trained on the boltuix/conll2025-ner dataset, it excels in information extraction, chatbots, search enhancement, and knowledge graph construction.

EntityBERT brings precise NER to the edge, enabling efficient AI for resource-constrained environments.

Key Features

• Lightweight: ~15 MB footprint, ~4.4M parameters.
• High Accuracy: 0.85 F1 score on conll2025-ner test set.
• 36 Entity Types: Supports 18 categories with BIO tagging.
• Edge-Optimized: Real-time inference with <30ms latency on Raspberry Pi.
• Offline Ready: No internet required.
• Versatile: Fine-tunable for other tasks like QA or intent classification.

Use Cases

• Information Extraction: Extract entities (👤 PERSON, 🌍 GPE, 🗓️ DATE) from news or reports.
• Chatbots: Enhance query understanding with entity recognition.
• Search Enhancement: Enable semantic search (e.g., “news about Paris in 2025”).
• Knowledge Graphs: Build structured connections between 🏢 ORG and 👤 PERSON.
• Domain-Specific NER: Fine-tune for medical 🩺, legal 📜, or financial 💸 applications.
• IoT Devices: Local entity extraction for smart home or wearable apps.

Supported Entity Labels

EntityBERT uses the BIO scheme, supporting 36 tags (18 entity types + O):

Tag	Purpose	Example
O	Outside entity	the, is
B-PERSON	Beginning of person	Elon
I-PERSON	Inside person	Musk
B-ORG	Beginning of organization	Tesla
I-ORG	Inside organization	Inc
B-GPE	Geopolitical entity	California
B-DATE	Beginning of date	March
I-DATE	Inside date	2025
B-LOC	Non-GPE location	Pacific
B-MONEY	Monetary value	$100
B-EVENT	Event	Olympics
B-FAC	Facility	Eiffel Tower
B-PRODUCT	Product	iPhone
B-LAW	Legal document	Constitution
B-NORP	Nationality/group	Democrat
B-CARDINAL	Cardinal number	1000
B-ORDINAL	Ordinal number	first
B-PERCENT	Percentage	50%
B-QUANTITY	Quantity	two liters
B-TIME	Time	noon
B-WORK_OF_ART	Work of art	Mona Lisa
B-LANGUAGE	Language	Spanish

Getting Started

Inference Example

from transformers import AutoTokenizer, AutoModelForTokenClassification
import torch

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("boltuix/EntityBERT")
model = AutoModelForTokenClassification.from_pretrained("boltuix/EntityBERT")

# Input text
text = "Elon Musk launched Tesla in California on March 2025."
inputs = tokenizer(text, return_tensors="pt")

# Run inference
with torch.no_grad():
    outputs = model(**inputs)
predictions = outputs.logits.argmax(dim=-1)

# Map predictions to labels
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
label_map = model.config.id2label
labels = [label_map[p.item()] for p in predictions[0]]

# Print results
for token, label in zip(tokens, labels):
    if token not in tokenizer.all_special_tokens:
        print(f"{token:15} → {label}")
                        Copy Code

Output:

Elon            → B-PERSON
Musk            → I-PERSON
launched        → O
Tesla           → B-ORG
in              → O
California      → B-GPE
on              → O
March           → B-DATE
2025            → I-DATE
.               → O
                        

Installation

pip install transformers torch pandas pyarrow seqeval
                        Copy Code

Requires Python 3.8+, ~15 MB for model, ~6.38 MB for dataset.

Performance Metrics

Evaluated on the conll2025-ner test set (~12,217 examples):

Metric	Score
Precision	0.84
Recall	0.86
F1 Score	0.85
Accuracy	0.91

Training Setup

• Hardware: NVIDIA GPU
• Training Time: ~1.5 hours
• Parameters: ~4.4M
• Optimizer: AdamW
• Learning Rate: 2e-5
• Batch Size: 16

Fine-Tuning Guide

Fine-tune EntityBERT on custom NER datasets:

# Install dependencies
!pip install evaluate transformers pandas datasets seqeval -q
import os
os.environ["WANDB_MODE"] = "disabled"

from transformers import BertTokenizerFast, AutoModelForTokenClassification, TrainingArguments, Trainer, DataCollatorForTokenClassification
from datasets import Dataset
import pandas as pd
import evaluate
import numpy as np

# Load dataset
df = pd.read_parquet("conll2025_ner.parquet")
dataset = Dataset.from_pandas(df)

# Extract unique tags
all_tags = set(tag for tags in df["ner_tags"] for tag in tags)
unique_tags = sorted(list(all_tags))
tag2id = {tag: i for i, tag in enumerate(unique_tags)}
id2tag = {i: tag for tag, i in tag2id.items()}

# Convert tags to IDs
def convert_tags_to_ids(example):
    example["ner_tags"] = [tag2id[tag] for tag in example["ner_tags"]]
    return example

dataset = dataset.map(convert_tags_to_ids)

# Split dataset
dataset_dict = {
    "train": dataset.filter(lambda x: x["split"] == "train"),
    "validation": dataset.filter(lambda x: x["split"] == "validation"),
    "test": dataset.filter(lambda x: x["split"] == "test")
}
dataset = datasets.DatasetDict(dataset_dict)

# Initialize tokenizer and model
tokenizer = BertTokenizerFast.from_pretrained("boltuix/bert-mini")
model = AutoModelForTokenClassification.from_pretrained("boltuix/bert-mini", num_labels=len(unique_tags))

# Tokenize and align labels
def tokenize_and_align_labels(examples, label_all_tokens=True):
    tokenized_inputs = tokenizer(examples["tokens"], truncation=True, is_split_into_words=True)
    labels = []
    for i, label in enumerate(examples["ner_tags"]):
        word_ids = tokenized_inputs.word_ids(batch_index=i)
        previous_word_idx = None
        label_ids = []
        for word_idx in word_ids:
            if word_idx is None:
                label_ids.append(-100)
            elif word_idx != previous_word_idx:
                label_ids.append(label[word_idx])
            else:
                label_ids.append(label[word_idx] if label_all_tokens else -100)
            previous_word_idx = word_idx
        labels.append(label_ids)
    tokenized_inputs["labels"] = labels
    return tokenized_inputs

tokenized_datasets = dataset.map(tokenize_and_align_labels, batched=True)

# Training arguments
args = TrainingArguments(
    output_dir="./entitybert_finetuned",
    eval_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    num_train_epochs=1,
    weight_decay=0.01,
    report_to="none"
)

# Data collator
data_collator = DataCollatorForTokenClassification(tokenizer)

# Evaluation metric
metric = evaluate.load("seqeval")

def compute_metrics(eval_preds):
    pred_logits, labels = eval_preds
    pred_logits = np.argmax(pred_logits, axis=2)
    predictions = [[unique_tags[p] for (p, l) in zip(prediction, label) if l != -100] for prediction, label in zip(pred_logits, labels)]
    true_labels = [[unique_tags[l] for (p, l) in zip(prediction, label) if l != -100] for prediction, label in zip(pred_logits, labels)]
    results = metric.compute(predictions=predictions, references=true_labels)
    return {
        "precision": results["overall_precision"],
        "recall": results["overall_recall"],
        "f1": results["overall_f1"],
        "accuracy": results["overall_accuracy"]
    }

# Initialize trainer
trainer = Trainer(
    model,
    args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=data_collator,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics
)

# Train
trainer.train()

# Save model
model.save_pretrained("./entitybert_finetuned")
tokenizer.save_pretrained("./entitybert_finetuned")
                        Copy Code

Evaluation Code

from transformers import AutoTokenizer, AutoModelForTokenClassification
from seqeval.metrics import classification_report
import torch

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("boltuix/EntityBERT")
model = AutoModelForTokenClassification.from_pretrained("boltuix/EntityBERT")

# Test data
texts = ["Elon Musk launched Tesla in California on March 2025."]
true_labels = [["B-PERSON", "I-PERSON", "O", "B-ORG", "O", "B-GPE", "O", "B-DATE", "I-DATE", "O"]]

pred_labels = []
for text in texts:
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    predictions = outputs.logits.argmax(dim=-1)[0].cpu().numpy()
    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
    word_ids = inputs.word_ids(batch_index=0)
    word_preds = []
    previous_word_idx = None
    for idx, word_idx in enumerate(word_ids):
        if word_idx is None or word_idx == previous_word_idx:
            continue
        label = model.config.id2label[predictions[idx]]
        word_preds.append(label)
        previous_word_idx = word_idx
    pred_labels.append(word_preds)

# Evaluate
print("Predicted:", pred_labels)
print("True     :", true_labels)
print("\nEvaluation Report:\n")
print(classification_report(true_labels, pred_labels))
                        Copy Code

Visualizing Tag Distribution

Analyze the distribution of NER tags in the dataset:

import pandas as pd
from collections import Counter
import matplotlib.pyplot as plt

# Load dataset
df = pd.read_parquet("conll2025_ner.parquet")
all_tags = [tag for tags in df["ner_tags"] for tag in tags]
tag_counts = Counter(all_tags)

# Plot
plt.figure(figsize=(12, 7))
plt.bar(tag_counts.keys(), tag_counts.values(), color="#36A2EB")
plt.title("CoNLL 2025 NER: Tag Distribution", fontsize=16)
plt.xlabel("NER Tag", fontsize=12)
plt.ylabel("Count", fontsize=12)
plt.xticks(rotation=45, ha="right", fontsize=10)
plt.grid(axis="y", linestyle="--", alpha=0.7)
plt.tight_layout()
plt.savefig("ner_tag_distribution.png")
plt.show()
                        Copy Code

Comparison to Other Models

Model	Dataset	Parameters	F1 Score	Size
EntityBERT	conll2025-ner	~4.4M	0.85	~15 MB
NeuroBERT-NER	conll2025-ner	~11M	0.86	~50 MB
BERT-base-NER	CoNLL-2003	~110M	~0.89	~400 MB
DistilBERT-NER	CoNLL-2003	~66M	~0.85	~200 MB

Frequently Asked Questions (FAQ)

Other Capabilities

While EntityBERT is optimized for NER, its bert-mini base supports fine-tuning for tasks like:

• Question Answering: Extract answers from text (e.g., “Who founded Tesla?” → “Elon Musk”).
• Intent Classification: Classify user intents (e.g., “Play music” → Play).
• Sentiment Analysis: Detect positive/negative sentiment.
• Semantic Similarity: Measure text similarity for clustering.

Dataset Details

• Entries: 143,709
• Size: 6.38 MB (Parquet)
• Splits: Train (~115,812), Validation (~15,680), Test (~12,217)
• Source: News, user-generated content, research corpora

License

Apache-2.0 License: Free to use. See LICENSE.

Support & Community

• Hugging Face Model Page
• Transformers Documentation
• Contact: boltuix@gmail.com

Conclusion

EntityBERT delivers precise, lightweight NER for edge AI, IoT, and mobile applications. With 36 entity types and a 0.85 F1 score, it’s perfect for information extraction, chatbots, and more. Explore it on Hugging Face!