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