This model is built by fine-tuning
SinBERT-large with an additional
Conditional Random Field (CRF) layer to ensure structurally valid predictions.
The model recognizes 4 entity types using the BIO tagging scheme (9 classes total including O):
1import tensorflow as tf
2from transformers import AutoTokenizer, TFAutoModel
3from tf2crf import CRF
4
5# 1. Define the custom model architecture
6class SinBERTCRFModel(tf.keras.Model):
7 def __init__(self, model_path, num_tags, **kwargs):
8 super(SinBERTCRFModel, self).__init__(**kwargs)
9 self.bert = TFAutoModel.from_pretrained(model_path, from_pt=False)
10 self.dropout = tf.keras.layers.Dropout(0.1)
11 self.classifier = tf.keras.layers.Dense(num_tags, name="classifier")
12 self.crf = CRF(units=num_tags, name="crf_layer")
13
14 def call(self, inputs, training=False):
15 outputs = self.bert(inputs['input_ids'], attention_mask=inputs['attention_mask'], training=training)
16 sequence_output = self.dropout(outputs.last_hidden_state, training=training)
17 emissions = self.classifier(sequence_output)
18 return self.crf(emissions)
19
20# 2. Initialize Tokenizer and Model
21model_id = "OmeshInusha999/SinBERT-NER-CRF"
22tokenizer = AutoTokenizer.from_pretrained(model_id, add_prefix_space=True)
23
24# Build the model architecture
25NUM_TAGS = 9
26model = SinBERTCRFModel("NLPC-UOM/SinBERT-large", NUM_TAGS)
27
28# Initialize model weights with a dummy forward pass
29dummy_input = {
30 "input_ids": tf.zeros((1, 128), dtype=tf.int32),
31 "attention_mask": tf.zeros((1, 128), dtype=tf.int32),
32}
33_ = model(dummy_input, training=False)
34
35# Load the trained CRF weights (Ensure you downloaded the checkpoint files from the repo)
36model.load_weights(f"./crf_model_checkpoint")
37
38# 3. Predict
39sentence = "අධික වර්ෂාවත් සමඟ හැටන් - කොළඹ ප්රධාන මාර්ගය අවදානමකට ලක්ව ඇත."
40raw_words = sentence.split()
41
42inputs = tokenizer(
43 raw_words,
44 is_split_into_words=True,
45 return_tensors="tf",
46 truncation=True,
47 max_length=128,
48 padding="max_length",
49)
50
51viterbi_sequence, _, _, _ = model({
52 "input_ids": inputs["input_ids"],
53 "attention_mask": inputs["attention_mask"]
54}, training=False)
55
56predictions = viterbi_sequence.numpy()[0]
57# Map predictions back to labels using your id2label mapping...