Named Entity Recognition (NER) with Deep Learning

A comprehensive Named Entity Recognition project that implements multiple BERT-based models and baseline approaches for biomedical entity extraction using BIO tagging scheme.

Project Overview

This project trains and evaluates multiple NER models:

• Baseline Model: RoBERTa model
• Main Models: BioBERT, and BioMedBERT models with CRF layers
• Task: Extract medical entities from biomedical text using BIO (Begin-Inside-Outside) tagging

Project Structure

NLP/
├── src/
│   ├── custom_model.py         # Custom BERT+CRF model implementation
│   ├── utils.py                # Utility functions for data processing and evaluation
│   └── __init__.py
├── data/
│   ├── mts_data.xml            # Original MTS biomedical data in BioC XML format
│   ├── train.csv               # Training data
│   ├── validation.csv          # Validation data
│   └── test.csv                # Test data
├── data_iob2/
│   ├── train.iob2              # Training data in IOB2 format
│   ├── validation.iob2         # Validation data in IOB2 format
│   ├── test.iob2               # Test data in IOB2 format
│   └── mts_data.iob2           # MTS data in IOB2 format
├── baseline/
│   ├── BiLSTM_baseline.ipynb   # Baseline model used during project proposal 
│   ├── span_f1.py              # Span-level F1 score calculation
│   └── [test files]            # Baseline test data in IOB2 format
├── results/
│   ├── MTS_results_*.csv       # Model results on MTS data
│   ├── MM_test_results_*.csv   # Model results on MM test data
│   └── [various result files]
├── main_train.py               # Main training pipeline for BERT-based models
├── baseline_train.py           # Training pipeline for baseline RoBERTa model
├── bio_tagger.py               # BIO tagging implementation
├── encode_tags.py              # Tag encoding utilities
├── convert_xml_to_iob2.ipynb   # Notebook to convert XML to IOB2 format
├── analysis.ipynb              # Analysis and evaluation notebook
├── requirements.txt            # Project dependencies
└── README.md                   # This file

Dependencies

Install required packages:

pip install -r requirements.txt

Key Libraries:

• torch: Deep learning framework
• transformers: Hugging Face transformers for BERT models
• pytorch_crf: CRF layer implementation
• seqeval: Evaluation metrics for NER tasks
• datasets: HuggingFace datasets library
• bioc: BioC XML format parsing
• pandas/numpy: Data manipulation
• matplotlib/seaborn: Visualization

Setup Instructions

1. Install Dependencies

pip install -r requirements.txt

2. Prepare Data

The data should be in IOB2 format (provided in data_iob2/ folder). If you have raw XML data:

• Use the convert_xml_to_iob2.ipynb notebook to convert BioC XML format to IOB2 format
• The conversion preserves entity annotations with BIO tagging scheme

3. Verify File Structure

Ensure all required files exist:

# Check data files
ls data_iob2/train.iob2
ls data_iob2/validation.iob2
ls data_iob2/test.iob2

Running the Models

Option 1: Baseline Model (RoBERTa + BiLSTM)

Run the baseline training:

python baseline_train.py

What it does:

• Loads training/validation/test data from data_iob2/ folder
• Trains a RoBERTa model for token classification
• Performs early stopping based on validation F1 score
• Saves results to results/ folder

Expected output:

• MTS_results_baseline.csv - Results on MTS test set
• MM_test_results_baseline.csv - Results on MM test set
• Model checkpoints in model cache

Option 2: Main Models (BERT with CRF)

Run the main training pipeline:

python main_train.py

What it does:

• Through a single main() function, the program takes the following chosen models as input to train it:
  1. BioBERT (dmis-lab/biobert-base-cased-v1.1)
  2. BioMedBERT (d4data/biomedical-bert-base-cased)
• Uses weighted CRF loss for better entity recognition
• Implements early stopping and model checkpointing
• Saves results to results/ folder

Expected output:

• MTS_results_biobert.csv - BioBERT results on MTS data
• MTS_results_biomed.csv - BioMedBERT results on MTS data
• Model checkpoints and evaluation metrics

Option 3: Using Google Colab

Training notebooks are available on Google Colab for GPU acceleration:

• Baseline Training: Colab Notebook
• Main Models Training: Colab Notebook

Data Format

IOB2 Format

The IOB2 format is a standard for NER tasks:

word    tag
Paris   B-LOCATION
is      O
nice    O
.       O

Each sentence is separated by an empty line. Tags follow the BIO scheme:

• B-ENTITY: Beginning of an entity
• I-ENTITY: Inside an entity (continuation)
• O: Outside any entity

Entity Types Supported:

• PERSON: Person names
• LOCATION: Geographic locations
• ORG: Organizations
• MEDICAL: Symptoms, History, Action
• O: No entity

Key Code Components

src/custom_model.py

BertCRFForNER: Custom model combining BERT with CRF layer

• BERT encoder for contextual representations
• Linear layer for emission scores
• CRF layer for sequence-level constraints
• Handles subword tokenization and label alignment

WeightedTrainer: Custom trainer class extending HuggingFace Trainer

• Implements weighted loss computation
• Proper state dict handling for CRF modules

src/utils.py

Key functions:

• read_file(): Parse IOB2 format files
• tokenize_and_align_labels(): Align subword tokens with original labels
• compute_metrics_model(): Calculate F1, precision, recall at token and entity level
• get_predictions(): Generate predictions on test data
• validate_bio() / fix_bio(): Validate and correct BIO tag sequences
• label_config(): Configure entity labels and their mappings

bio_tagger.py

BIOTagger: Converts character-level entity annotations to BIO-tagged sequences

• Handles entity boundary detection
• Converts between different entity type mappings
• Produces valid BIO sequences

Training Pipeline Details

1. Data Loading

train_data = read_file('data_iob2/train.iob2')  # Returns list of (words, tags) tuples

2. Tokenization & Label Alignment

• Converts words to subword tokens using model tokenizer
• Aligns entity labels to subword tokens (first token gets label, rest ignored)
• Prepares features for model input

3. Model Training

• Uses Trainer API from HuggingFace transformers
• Early stopping when validation F1 plateaus
• Saves best model checkpoint
• Computes metrics: F1 score, precision, recall (both token and entity-level)

4. Evaluation

• Generates predictions on test set
• Creates classification reports per entity type
• Saves results to CSV files

Results Interpretation

Results files contain metrics for each entity type:

• Precision: Percentage of predicted entities that are correct
• Recall: Percentage of gold entities that were found
• F1 Score: Harmonic mean of precision and recall
• Support: Number of entities in test set

Hyperparameters

Common settings in training scripts:

• Batch size: 8-16 (adjust based on GPU memory)
• Learning rate: 2e-5 (standard for BERT fine-tuning)
• Epochs: 20 (with early stopping)
• Random seed: 52 (for reproducibility)
• Max sequence length: 512 (BERT max)

GPU Requirements

• Recommended: NVIDIA GPU with 8GB+ VRAM
• CPU: Significantly slower
• Colab: Use GPU runtime for faster training

References

• Transformer Models
• BIO Tagging Scheme
• SeqEval Metrics
• PyTorch CRF

License

This is an academic project for Natural Language Processing course.