| # Custom BERT NER Model |
|
|
| This repository contains a BERT-based Named Entity Recognition (NER) model fine-tuned on the CoNLL-2003 dataset. The model is trained to identify common named entity types such as persons, organizations, locations, and miscellaneous entities. |
|
|
| --- |
|
|
| ## Model Details |
|
|
| - **Model architecture:** BERT (bert-base-cased) |
| - **Task:** Token classification / Named Entity Recognition (NER) |
| - **Training data:** CoNLL-2003 dataset (~14,000 training samples) |
| - **Number of epochs:** 5 |
| - **Framework:** Hugging Face Transformers + Datasets |
| - **Device:** CUDA-enabled GPU for training and inference |
| - **WandB:** Disabled during training |
|
|
| --- |
|
|
| ## Usage |
|
|
| You can use this model for token classification to identify named entities in your text. |
|
|
| ### Installation |
|
|
| ```python |
| pip install transformers datasets torch |
| ``` |
|
|
| ## Load the model and tokenizer |
|
|
| ```pyhton |
| |
| from transformers import BertTokenizerFast, BertForTokenClassification |
| import torch |
| |
| model_name_or_path = "AventIQ-AI/Custom-BERT-NER-Model" |
| |
| tokenizer = BertTokenizerFast.from_pretrained(model_name_or_path) |
| model = BertForTokenClassification.from_pretrained(model_name_or_path) |
| |
| model.to("cuda") # or "cpu" |
| model.eval() |
| ``` |
|
|
| ## Example inference |
|
|
| ```python |
| |
| text = "Hi, I am Deepak and I am living in Delhi." |
| |
| tokens = tokenizer(text, return_tensors="pt").to(model.device) |
| outputs = model(**tokens) |
| predictions = torch.argmax(outputs.logits, dim=2) |
| |
| labels = [model.config.id2label[p.item()] for p in predictions[0]] |
| for token, label in zip(tokenizer.tokenize(text), labels): |
| print(f"{token}: {label}") |
| ``` |
|
|
| ## Training Details |
|
|
| - Dataset: CoNLL-2003, loaded via the Hugging Face datasets library |
|
|
| - Optimizer: AdamW |
|
|
| - Learning Rate: 5e-5 |
|
|
| - Batch Size: 16 |
|
|
| - Max Sequence Length: 128 |
|
|
| - Epochs: 5 |
|
|
| - Evaluation: Performed on validation split (if applicable) |
|
|
| - Quantization: Applied post-training for model size reduction (optional) |
|
|
| ## Limitations |
|
|
| - The model may not generalize well to unseen entity types or domains outside CoNLL-2003. |
|
|
| - It can occasionally mislabel entities, especially for rare or new names. |
|
|
| - A CUDA-enabled GPU is required for efficient training and inference. |
|
|
|
|