SinBERT-NER-CRF

SinBERT-NER-CRF: Sinhala Named Entity Recognition Model

A state-of-the-art Named Entity Recognition (NER) tagger for the Sinhala (සිංහල) language.

Note: This is a specialized token classification model (tagger), not a Generative Large Language Model (LLM). It is designed strictly for extracting structured entities from text.

This model is built by fine-tuning SinBERT-large with an additional Conditional Random Field (CRF) layer to ensure structurally valid predictions.

Model Architecture & Technical Details

Why use a CRF Layer?

Standard transformer token classifiers predict the tag for each word independently. This can lead to structurally impossible sequences (e.g., an I-PER tag following a B-LOC tag).

By adding a Conditional Random Field (CRF) layer on top of the transformer's emissions, the model learns the transition probabilities between different tags. It decodes the sequence globally (using the Viterbi algorithm) to find the most probable chain of tags, effectively eliminating invalid "BIO" tag transitions and improving boundary detection.

Supported Entities

The model recognizes 4 entity types using the BIO tagging scheme (9 classes total including O):

Evaluation Results

The model was evaluated on the test split (384 samples, 861 entities) of the polyglots/Sinhala-NER dataset. Adding the CRF layer resulted in a significant improvement across all metrics compared to the base Transformer model.

Overall Performance (Entity-Level)

Per-Entity Performance (CRF Model)

Usage

Because this model uses a custom CRF layer via tf2crf, it cannot be loaded directly using the standard TFAutoModelForTokenClassification.from_pretrained() pipeline. You must rebuild the architecture and load the weights.

Requirements:

pip install tensorflow==2.15.0 transformers tf2crf

Inference Code:

import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModel
from tf2crf import CRF

# 1. Define the custom model architecture
class SinBERTCRFModel(tf.keras.Model):
    def __init__(self, model_path, num_tags, **kwargs):
        super(SinBERTCRFModel, self).__init__(**kwargs)
        self.bert = TFAutoModel.from_pretrained(model_path, from_pt=False)
        self.dropout = tf.keras.layers.Dropout(0.1)
        self.classifier = tf.keras.layers.Dense(num_tags, name="classifier")
        self.crf = CRF(units=num_tags, name="crf_layer")

    def call(self, inputs, training=False):
        outputs = self.bert(inputs['input_ids'], attention_mask=inputs['attention_mask'], training=training)
        sequence_output = self.dropout(outputs.last_hidden_state, training=training)
        emissions = self.classifier(sequence_output)
        return self.crf(emissions)

# 2. Initialize Tokenizer and Model
model_id = "OmeshInusha999/SinBERT-NER-CRF" 
tokenizer = AutoTokenizer.from_pretrained(model_id, add_prefix_space=True)

# Build the model architecture
NUM_TAGS = 9
model = SinBERTCRFModel("NLPC-UOM/SinBERT-large", NUM_TAGS)

# Initialize model weights with a dummy forward pass
dummy_input = {
    "input_ids": tf.zeros((1, 128), dtype=tf.int32),
    "attention_mask": tf.zeros((1, 128), dtype=tf.int32),
}
_ = model(dummy_input, training=False)

# Load the trained CRF weights (Ensure you downloaded the checkpoint files from the repo)
model.load_weights(f"./crf_model_checkpoint")

# 3. Predict
sentence = "අධික වර්ෂාවත් සමඟ හැටන් - කොළඹ ප්‍රධාන මාර්ගය අවදානමකට ලක්ව ඇත."
raw_words = sentence.split()

inputs = tokenizer(
    raw_words,
    is_split_into_words=True,
    return_tensors="tf",
    truncation=True,
    max_length=128,
    padding="max_length",
)

viterbi_sequence, _, _, _ = model({
    "input_ids": inputs["input_ids"],
    "attention_mask": inputs["attention_mask"]
}, training=False)

predictions = viterbi_sequence.numpy()[0]
# Map predictions back to labels using your id2label mapping...

Limitations

• Optimized specifically for formal/news Sinhala text.
• Maximum sequence length is limited to 128 tokens.
• Cannot be used directly with HuggingFace pipeline() due to the custom TensorFlow CRF layer.

Acknowledgements

• Base Model: NLPC-UOM/SinBERT-large
• Dataset: polyglots/Sinhala-NER