ILoveBacteria/bilstm-sequence-labeler

BiLSTM Sequence Labeler (NER, POS, Chunking)

This repository hosts a custom Bidirectional LSTM (BiLSTM) sequence labeling model trained on the standard CoNLL-2003 dataset. The network uses a shared embedding layer and a bidirectional recurrent backbone to perform three tasks simultaneously:

1. Named Entity Recognition (NER)
2. Part-of-Speech Tagging (POS)
3. Chunk Tagging

For full training details, notebooks, and source files, visit the GitHub Repository.

Model Architecture

• Embedding Layer: 128-dimensional lookup table.
• Recurrent Layer: Bidirectional LSTM (hidden_dim=128, yielding a 256-dimensional concatenated sequence representation).
• Classification Heads: Three independent parallel linear layers map the shared hidden states directly to class distributions for NER, POS, and Chunk tags.
• Padding Handle: Masked sequence computation utilizing ignore_index=-100 ensuring pad tokens do not affect loss or performance evaluations.

---

Training and Evaluation Curves

Plot Breakdown:

• Loss Convergence: Illustrates the trajectory of the optimization objective across training epochs. The training loss displays a consistent downward gradient, reflecting successful gradient updates, while the evaluation loss tracking curve confirms stable generalization without premature over-fitting or cross-entropy divergence.
• Accuracy Metrics: Tracks token-level micro/macro-accuracy metrics across epochs. The validation tracking demonstrates a swift initial trajectory before stabilizing into an asymptotic plateau, indicating that the shared bidirectional recurrent backbone effectively captured optimal contextual representations for joint task prediction early in the training phase.

---

Evaluation Results (NER Task)

The model achieved an overall Accuracy of 94% on the evaluation set. Below is the detailed token-level classification report across all 9 target categories:

Class / Tag ID	Precision	Recall	F1-Score	Support
0 (e.g., O)	0.96	0.99	0.98	81,082
1	0.89	0.73	0.80	3,459
2	0.90	0.78	0.83	2,463
3	0.83	0.66	0.74	3,002
4	0.81	0.69	0.75	1,586
5	0.84	0.78	0.81	3,505
6	0.62	0.68	0.65	514
7	0.82	0.68	0.74	1,624
8	0.66	0.63	0.65	562
Macro Average	0.82	0.74	0.77	97,797
Weighted Average	0.94	0.94	0.94	97,797

---

Sample Visualization (displaCy)

Visualization Breakdown:

• Token Alignment: Demonstrates the model's live out-of-sample performance on an unseen text sequence. The visualization shows tokens highlighted and labeled dynamically according to the extracted entity types (such as ORG for organizations, LOC for locations, or PER for persons).
• Multi-Task Ingestion: While displacy.render specifically highlights the structural Named Entity Recognition (NER) tag segments extracted by the token prediction head, these entity boundaries are intrinsically anchored by the shared hidden representations refined by the parallel POS and Chunk tagging heads during the joint optimization process.

---

How to Use

Since this model implements a completely custom neural architecture, make sure to pass trust_remote_code=True when loading via AutoModel.

import torch
import json
import urllib.request
from transformers import AutoConfig, AutoModel

# 1. Define repository identifier and vocabulary URL
repo_id = "ILoveBacteria/bilstm-sequence-labeler"
vocab_url = "https://huggingface.co/ILoveBacteria/bilstm-sequence-labeler/resolve/main/vocab.json"

# 2. Load custom model architecture and weights
model = AutoModel.from_pretrained(repo_id, trust_remote_code=True)
model.eval()

# 3. Download and load the vocabulary file directly from the link
with urllib.request.urlopen(vocab_url) as response:
    vocab = json.loads(response.read().decode("utf-8"))

token2id = vocab["token2id"]

# Simple inference example
text = "EU rejects German call to boycott British lamb ."
tokens = text.split()
input_ids = [token2id.get(token, token2id["<unk>"]) for token in tokens]
input_tensor = torch.tensor([input_ids]) # Add batch dimension

with torch.no_grad():
    outputs = model(input_tensor)

# Extract predictions for tasks
ner_predictions = outputs["ner"].argmax(dim=-1)[0]
print("Predicted NER IDs:", ner_predictions.tolist())

Advanced Visualization (displaCy)

If you are running inference within a Jupyter Notebook or Google Colab, you can use the script below to generate a beautiful, interactive visual layout of the predicted Named Entities using SpaCy's displacy engine.

from spacy import displacy

# 1. Map for CoNLL-2003 NER indices to human-readable labels
ner_labels = {
    0: "O", 1: "B-PER", 2: "I-PER", 3: "B-ORG", 4: "I-ORG",
    5: "B-LOC", 6: "I-LOC", 7: "B-MISC", 8: "I-MISC"
}

# 2. Convert predicted tensor to a flat NumPy array or list
preds = ner_predictions.cpu().numpy()

# 3. Reconstruct the string text and track exact character offsets for displaCy
text = " ".join(tokens)
ents = []
char_offsets = []
cursor = 0

for tok in tokens:
    char_offsets.append((cursor, cursor + len(tok)))
    cursor += len(tok) + 1  # +1 accounts for the space between tokens

# 4. Construct the entity metadata required by displaCy
for i, pred_idx in enumerate(preds):
    label = ner_labels.get(int(pred_idx), "O")
    if label != "O":
        start, end = char_offsets[i]
        ents.append({"start": start, "end": end, "label": label})

# 5. Pack the document structural payload
doc_data = {
    "text": text,
    "ents": ents
}

# 6. Render the visualization inline (Jupyter Notebook / Google Colab)
displacy.render(doc_data, style="ent", manual=True, jupyter=True)