🧠 NERClassifier-BERT-WikiAnn
A BERT-based Named Entity Recognition (NER) model fine-tuned on the WikiAnn English dataset. It classifies tokens into entity types like Person (PER), Location (LOC), and Organization (ORG). This model is suitable for applications like document tagging, resume parsing, and chatbots.
✨ Model Highlights
- 📌 Based on bert-base-cased
- 🔍 Fine-tuned on the WikiAnn (en) NER dataset
- ⚡ Supports prediction of 3 core entity types: PER, LOC, ORG
- 💾 Lightweight and compatible with both CPU and GPU inference environment
🧠 Intended Uses
- ✅ Resume and document parsing
- ✅ News article analysis
- ✅ Question answering pipelines
- ✅ Chatbots and virtual assistants
- ✅ Information retrieval and tagging
🚫 Limitations
❌ Trained on English-only Wiki-based text
❌ Performance may degrade on informal or non-English texts
❌ Not designed for nested or overlapping entities
❌ Accuracy may drop on very long sequences (>128 tokens)
🏋️♂️ Training Details
| Field | Value |
|---|---|
| Base Model | bert-base-cased |
| Dataset | WikiAnn (English) |
| Framework | PyTorch with 🤗 Transformers |
| Epochs | 3 |
| Batch Size | 16 |
| Max Length | 128 tokens |
| Optimizer | AdamW |
| Loss | CrossEntropyLoss (token-level) |
| Device | Trained on CUDA-enabled GPU |
📊 Evaluation Metrics
| Metric | Score |
|---|---|
| Accuracy | 0.92 |
| F1-Score | 0.92 |
| Precision | 0.92 |
| Recall | 0.92 |
🔎 Label Mapping
| Label ID | Entity Type |
|---|---|
| 0 | O |
| 1 | B-PER |
| 2 | I-PER |
| 3 | B-ORG |
| 4 | I-ORG |
| 5 | B-LOC |
| 6 | I-LOC |
🚀 Usage
from transformers import BertTokenizerFast, BertForTokenClassification
from transformers import pipeline
import torch
model_name = "AventIQ-AI/NER-AI-wikiann-model"
tokenizer = BertTokenizerFast.from_pretrained(model_name)
model = BertForTokenClassification.from_pretrained(model_name)
model.eval()
#Labelling
label_list = dataset["train"].features["ner_tags"].feature.names
model.config.id2label = {i: label for i, label in enumerate(label_list)}
model.config.label2id = {label: i for i, label in enumerate(label_list)}
ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple")
#Inference
ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple")
test_sentence = "Bill Gates is the CEO of Microsoft and lives in the United States."
ner_results = ner_pipeline(test_sentence)
print("\n📌 Inference Results:")
for entity in ner_results:
print(f"Entity: {entity['word']}\tType: {entity['entity_group']}\tConfidence: {entity['score']:.3f}")
# Test example
print("Bill Gates is the CEO of Microsoft and lives in the United States.")
- 🧩 Quantization
- Post-training static quantization applied using PyTorch to reduce model size and accelerate inference on edge devices.
🗂 Repository Structure
.
├── model/ # Quantized model files
├── tokenizer_config/ # Tokenizer and vocab files
├── model.safensors/ # Fine-tuned model in safetensors format
├── README.md # Model card
🤝 Contributing
Open to improvements and feedback! Feel free to submit a pull request or open an issue if you find any bugs or want to enhance the model.
- Downloads last month
- 5