1# 🛠️ Step 1: Install required libraries quietly
2!pip install evaluate transformers datasets tokenizers seqeval pandas pyarrow -q
3
4# 🚫 Step 2: Disable Weights & Biases (WandB)
5import os
6os.environ["WANDB_MODE"] = "disabled"
7
8# 📚 Step 2: Import necessary libraries
9import pandas as pd
10import datasets
11import numpy as np
12from transformers import BertTokenizerFast
13from transformers import DataCollatorForTokenClassification
14from transformers import AutoModelForTokenClassification
15from transformers import TrainingArguments, Trainer
16import evaluate
17from transformers import pipeline
18from collections import defaultdict
19import json
20
21# 📥 Step 3: Load the CoNLL-2025 NER dataset from Parquet
22# Download : https://huggingface.co/datasets/boltuix/conll2025-ner/blob/main/conll2025_ner.parquet
23parquet_file = "conll2025_ner.parquet"
24df = pd.read_parquet(parquet_file)
25
26# 🔍 Step 4: Convert pandas DataFrame to Hugging Face Dataset
27conll2025 = datasets.Dataset.from_pandas(df)
28
29# 🔎 Step 5: Inspect the dataset structure
30print("Dataset structure:", conll2025)
31print("Dataset features:", conll2025.features)
32print("First example:", conll2025[0])
33
34# 🏷️ Step 6: Extract unique tags and create mappings
35# Since ner_tags are strings, collect all unique tags
36all_tags = set()
37for example in conll2025:
38 all_tags.update(example["ner_tags"])
39unique_tags = sorted(list(all_tags)) # Sort for consistency
40num_tags = len(unique_tags)
41tag2id = {tag: i for i, tag in enumerate(unique_tags)}
42id2tag = {i: tag for i, tag in enumerate(unique_tags)}
43print("Number of unique tags:", num_tags)
44print("Unique tags:", unique_tags)
45
46# 🔧 Step 7: Convert string ner_tags to indices
47def convert_tags_to_ids(example):
48 example["ner_tags"] = [tag2id[tag] for tag in example["ner_tags"]]
49 return example
50
51conll2025 = conll2025.map(convert_tags_to_ids)
52
53# 📊 Step 8: Split dataset based on 'split' column
54dataset_dict = {
55 "train": conll2025.filter(lambda x: x["split"] == "train"),
56 "validation": conll2025.filter(lambda x: x["split"] == "validation"),
57 "test": conll2025.filter(lambda x: x["split"] == "test")
58}
59conll2025 = datasets.DatasetDict(dataset_dict)
60print("Split dataset structure:", conll2025)
61
62# 🪙 Step 9: Initialize the tokenizer
63tokenizer = BertTokenizerFast.from_pretrained("boltuix/bert-mini")
64
65# 📝 Step 10: Tokenize an example text and inspect
66example_text = conll2025["train"][0]
67tokenized_input = tokenizer(example_text["tokens"], is_split_into_words=True)
68tokens = tokenizer.convert_ids_to_tokens(tokenized_input["input_ids"])
69word_ids = tokenized_input.word_ids()
70print("Word IDs:", word_ids)
71print("Tokenized input:", tokenized_input)
72print("Length of ner_tags vs input IDs:", len(example_text["ner_tags"]), len(tokenized_input["input_ids"]))
73
74# 🔄 Step 11: Define function to tokenize and align labels
75def tokenize_and_align_labels(examples, label_all_tokens=True):
76 """
77 Tokenize inputs and align labels for NER tasks.
78
79 Args:
80 examples (dict): Dictionary with tokens and ner_tags.
81 label_all_tokens (bool): Whether to label all subword tokens.
82
83 Returns:
84 dict: Tokenized inputs with aligned labels.
85 """
86 tokenized_inputs = tokenizer(examples["tokens"], truncation=True, is_split_into_words=True)
87 labels = []
88 for i, label in enumerate(examples["ner_tags"]):
89 word_ids = tokenized_inputs.word_ids(batch_index=i)
90 previous_word_idx = None
91 label_ids = []
92 for word_idx in word_ids:
93 if word_idx is None:
94 label_ids.append(-100) # Special tokens get -100
95 elif word_idx != previous_word_idx:
96 label_ids.append(label[word_idx]) # First token of word gets label
97 else:
98 label_ids.append(label[word_idx] if label_all_tokens else -100) # Subwords get label or -100
99 previous_word_idx = word_idx
100 labels.append(label_ids)
101 tokenized_inputs["labels"] = labels
102 return tokenized_inputs
103
104# 🧪 Step 12: Test the tokenization and label alignment
105q = tokenize_and_align_labels(conll2025["train"][0:1])
106print("Tokenized and aligned example:", q)
107
108# 📋 Step 13: Print tokens and their corresponding labels
109for token, label in zip(tokenizer.convert_ids_to_tokens(q["input_ids"][0]), q["labels"][0]):
110 print(f"{token:_<40} {label}")
111
112# 🔧 Step 14: Apply tokenization to the entire dataset
113tokenized_datasets = conll2025.map(tokenize_and_align_labels, batched=True)
114
115# 🤖 Step 15: Initialize the model with the correct number of labels
116model = AutoModelForTokenClassification.from_pretrained("boltuix/bert-mini", num_labels=num_tags)
117
118# ⚙️ Step 16: Set up training arguments
119args = TrainingArguments(
120 "boltuix/bert-ner",
121 eval_strategy="epoch", # Changed evaluation_strategy to eval_strategy
122 learning_rate=2e-5,
123 per_device_train_batch_size=16,
124 per_device_eval_batch_size=16,
125 num_train_epochs=1,
126 weight_decay=0.01,
127 report_to="none"
128)
129# 📊 Step 17: Initialize data collator for dynamic padding
130data_collator = DataCollatorForTokenClassification(tokenizer)
131
132# 📈 Step 18: Load evaluation metric
133metric = evaluate.load("seqeval")
134
135# 🏷️ Step 19: Set label list and test metric computation
136label_list = unique_tags
137print("Label list:", label_list)
138example = conll2025["train"][0]
139labels = [label_list[i] for i in example["ner_tags"]]
140print("Metric test:", metric.compute(predictions=[labels], references=[labels]))
141
142# 📉 Step 20: Define function to compute evaluation metrics
143def compute_metrics(eval_preds):
144 """
145 Compute precision, recall, F1, and accuracy for NER.
146
147 Args:
148 eval_preds (tuple): Predicted logits and true labels.
149
150 Returns:
151 dict: Evaluation metrics.
152 """
153 pred_logits, labels = eval_preds
154 pred_logits = np.argmax(pred_logits, axis=2)
155 predictions = [
156 [label_list[p] for (p, l) in zip(prediction, label) if l != -100]
157 for prediction, label in zip(pred_logits, labels)
158 ]
159 true_labels = [
160 [label_list[l] for (p, l) in zip(prediction, label) if l != -100]
161 for prediction, label in zip(pred_logits, labels)
162 ]
163 results = metric.compute(predictions=predictions, references=true_labels)
164 return {
165 "precision": results["overall_precision"],
166 "recall": results["overall_recall"],
167 "f1": results["overall_f1"],
168 "accuracy": results["overall_accuracy"],
169 }
170
171# 🚀 Step 21: Initialize and train the trainer
172trainer = Trainer(
173 model,
174 args,
175 train_dataset=tokenized_datasets["train"],
176 eval_dataset=tokenized_datasets["validation"],
177 data_collator=data_collator,
178 tokenizer=tokenizer,
179 compute_metrics=compute_metrics
180)
181trainer.train()
182
183# 💾 Step 22: Save the fine-tuned model
184model.save_pretrained("boltuix/bert-ner")
185tokenizer.save_pretrained("tokenizer")
186
187# 🔗 Step 23: Update model configuration with label mappings
188id2label = {str(i): label for i, label in enumerate(label_list)}
189label2id = {label: str(i) for i, label in enumerate(label_list)}
190config = json.load(open("boltuix/bert-ner/config.json"))
191config["id2label"] = id2label
192config["label2id"] = label2id
193json.dump(config, open("boltuix/bert-ner/config.json", "w"))
194
195# 🔄 Step 24: Load the fine-tuned model
196model_fine_tuned = AutoModelForTokenClassification.from_pretrained("boltuix/bert-ner")
197
198# 🛠️ Step 25: Create a pipeline for NER inference
199nlp = pipeline("token-classification", model=model_fine_tuned, tokenizer=tokenizer)
200
201# 📝 Step 26: Perform NER on an example sentence
202example = "On July 4th, 2023, President Joe Biden visited the United Nations headquarters in New York to deliver a speech about international law and donated $5 million to relief efforts."
203ner_results = nlp(example)
204print("NER results for first example:", ner_results)
205
206# 📍 Step 27: Perform NER on a property address and format output
207example = "This page contains information about the property located at 1275 Kinnear Rd, Columbus, OH, 43212."
208ner_results = nlp(example)
209
210# 🧹 Step 28: Process NER results into structured entities
211entities = defaultdict(list)
212current_entity = ""
213current_type = ""
214
215for item in ner_results:
216 entity = item["entity"]
217 word = item["word"]
218 if word.startswith("##"):
219 current_entity += word[2:] # Handle subword tokens
220 elif entity.startswith("B-"):
221 if current_entity and current_type:
222 entities[current_type].append(current_entity.strip())
223 current_type = entity[2:].lower()
224 current_entity = word
225 elif entity.startswith("I-") and entity[2:].lower() == current_type:
226 current_entity += " " + word # Continue same entity
227 else:
228 if current_entity and current_type:
229 entities[current_type].append(current_entity.strip())
230 current_entity = ""
231 current_type = ""
232
233# Append final entity if exists
234if current_entity and current_type:
235 entities[current_type].append(current_entity.strip())
236
237# 📤 Step 29: Output the final JSON
238final_json = dict(entities)
239print("Structured NER output:")
240print(json.dumps(final_json, indent=2))