KatherineFer/NER_DDI_SU

DrugNER with PubMedBERT on the DDI 2013 Corpus

This project builds a Named Entity Recognition (NER) model to identify drug-related entities in biomedical text using the [DDI 2013 Challenge Corpus] and fine-tunes the domain-specific PubMedBERT model.

Project Structure

DDICorpus/
├── Train/
│   ├── DrugBank/
│   └── MedLine/
├── Test/
│   └── Test for DrugNER task/
├── NER_DDI/
│   ├── parsed_data.json
│   ├── tokenized_dataset/
│   ├── val_dataset/
│   ├── models/
│   │   └── final_model/
│   └── eval_results/


## Model

- **Base model:** `microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract`
- **Fine-tuned on:** DDI 2013 training set
- **Entity types:** `drug`, `brand`, `group`, `drug_n`
- **Labels format:** BIO tagging scheme
- **Final evaluation:** on the DDI Test Set (unseen XML files)

## Training Pipeline

1. **Parse & label the XML corpus** (BIO format aligned to tokenizer offsets)
2. **Create HuggingFace Dataset** (`datasets.Dataset`)
3. **Tokenize with alignment**
4. **Fine-tune model using `Trainer`**
5. **Save model and tokenizer**
6. **Evaluate on test set**
7. **Compute metrics with `seqeval` + confusion matrix**

## Final Results (Test Set)

| Entity Type | Precision | Recall | F1-score |
|-------------|-----------|--------|----------|
| Drug        | 0.7820    | 0.8940 | 0.8342   |
| Brand       | 0.8780    | 0.6102 | 0.7200   |
| Group       | 0.8293    | 0.8774 | 0.8527   |
| Drug_n      | 0.4953    | 0.4417 | 0.4670   |
| **Overall** | **0.7553**| **0.7862** | **0.7704** |

## Example Inference

```python
sentence = "Aspirin and paracetamol were found to interact against each other."

entities = [
  {'type': 'drug', 'text': 'Aspirin'},
  {'type': 'drug', 'text': 'paracetamol'}
]