NeuroBERT-NER: Advanced Named Entity Recognition for Edge AI and IoT

Overview

NeuroBERT-NER is a fine-tuned transformer model for Named Entity Recognition (NER), built on boltuix/NeuroBERT-Mini. With ~11M parameters and a size of ~50 MB, it identifies 36 entity types (e.g., PERSON, ORG, GPE, DATE) in English text. Trained on the boltuix/conll2025-ner dataset, it’s optimized for edge AI, IoT devices, and mobile applications, excelling in information extraction, chatbots, search enhancement, and knowledge graphs.

NeuroBERT-NER delivers high-accuracy NER for edge devices, powering efficient AI applications.

Key Features

• Compact: ~50 MB footprint, ~11M parameters.
• High Accuracy: 0.86 F1 score on conll2025-ner test set.
• 36 Entity Types: Supports 18 categories with BIO tagging.
• Edge-Optimized: Real-time inference with low latency.
• Offline Ready: No internet required.
• Extensible: Fine-tunable for tasks like QA or intent classification.

Use Cases

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

Supported Entity Labels

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

Tag	Purpose	Example
O	Outside entity	the, is
B-PERSON	Beginning of person	Barack
I-PERSON	Inside person	Obama
B-ORG	Beginning of organization	Microsoft
I-ORG	Inside organization	Corp
B-GPE	Geopolitical entity	Seattle
B-DATE	Beginning of date	January
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/NeuroBERT-NER")
model = AutoModelForTokenClassification.from_pretrained("boltuix/NeuroBERT-NER")

# Input text
text = "Barack Obama visited Microsoft headquarters in Seattle on January 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:

Barack          → B-PERSON
Obama           → I-PERSON
visited         → O
Microsoft       → B-ORG
headquarters    → O
in              → O
Seattle         → B-GPE
on              → O
January         → B-DATE
2025            → I-DATE
.               → O
                        

Installation

pip install transformers torch pandas pyarrow seqeval
                        Copy Code

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

Performance Metrics

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

Metric	Score
Precision	0.85
Recall	0.87
F1 Score	0.86
Accuracy	0.92

Training Setup

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

Fine-Tuning Guide

Fine-tune NeuroBERT-NER on custom NER datasets:

# Install dependencies
!pip install transformers datasets tokenizers seqeval pandas pyarrow -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/NeuroBERT-Mini")
model = AutoModelForTokenClassification.from_pretrained("boltuix/NeuroBERT-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="./neurobert_ner_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("./neurobert_ner_finetuned")
tokenizer.save_pretrained("./neurobert_ner_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/NeuroBERT-NER")
model = AutoModelForTokenClassification.from_pretrained("boltuix/NeuroBERT-NER")

# Test data
texts = ["Barack Obama visited Microsoft in Seattle on January 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:

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
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
spaCy (en_core_web_lg)	OntoNotes	-	~0.83	~800 MB

Frequently Asked Questions (FAQ)

Other Capabilities

While optimized for NER, NeuroBERT-NER’s NeuroBERT-Mini base supports fine-tuning for:

• Question Answering: Answer questions from text (e.g., “Who visited Microsoft?” → “Barack Obama”).
• Intent Classification: Classify 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

Carbon Impact

Training emitted ~50g CO₂eq over ~2 hours on an NVIDIA GPU, estimated via ML Impact tool.

License

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

Support & Community

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

Conclusion

NeuroBERT-NER offers advanced NER with a 0.86 F1 score, optimized for edge AI and IoT. From chatbots to knowledge graphs, it’s a versatile solution for 2025. Explore it on Hugging Face!