Create training/train_ner.py

training/train_ner.py

@@ -0,0 +1,88 @@
+import json, argparse
+from datasets import Dataset
+from transformers import AutoTokenizer, AutoModelForTokenClassification, DataCollatorForTokenClassification, TrainingArguments, Trainer
+from training.utils import compute_metrics_ner
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--model_name", default="bert-base-cased")
+parser.add_argument("--train_json", required=True, help="JSONL with {'tokens': [...], 'ner_tags': [...]} per line")
+parser.add_argument("--eval_json", required=True)
+parser.add_argument("--text_col", default="tokens")
+parser.add_argument("--label_col", default="ner_tags")
+parser.add_argument("--labels_file", default="training/labels_ner.json")
+parser.add_argument("--output_dir", default="./outputs/ner")
+parser.add_argument("--epochs", type=int, default=5)
+parser.add_argument("--batch_size", type=int, default=8)
+parser.add_argument("--lr", type=float, default=3e-5)
+args = parser.parse_args()
+
+def load_jsonl(path):
+    rows = []
+    with open(path, "r", encoding="utf-8") as f:
+        for line in f:
+            rows.append(json.loads(line))
+    return rows
+
+train_rows = load_jsonl(args.train_json)
+eval_rows  = load_jsonl(args.eval_json)
+
+with open(args.labels_file, "r") as f:
+    label_list = json.load(f)  # e.g., ["O","B-ORG","I-ORG","B-MONEY","I-MONEY","B-DATE","I-DATE","B-TICKER","I-TICKER"]
+
+tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+
+def align_labels_with_tokens(tokens, labels):
+    # labels are per-token already; convert to ids
+    label2id = {l:i for i,l in enumerate(label_list)}
+    return [label2id[l] for l in labels]
+
+def encode_batch(batch):
+    tokenized = tokenizer(batch[args.text_col], is_split_into_words=True, truncation=True, padding=True)
+    encoded_labels = []
+    for i, labels in enumerate(batch[args.label_col]):
+        word_ids = tokenized.word_ids(batch_index=i)
+        label_ids = []
+        j = 0
+        for w_id in word_ids:
+            if w_id is None:
+                label_ids.append(-100)
+            else:
+                label_ids.append(align_labels_with_tokens(batch[args.text_col][i], labels)[w_id])
+        encoded_labels.append(label_ids)
+    tokenized["labels"] = encoded_labels
+    return tokenized
+
+train_ds = Dataset.from_list(train_rows).map(encode_batch, batched=True, remove_columns=[args.text_col, args.label_col])
+eval_ds  = Dataset.from_list(eval_rows).map(encode_batch,  batched=True, remove_columns=[args.text_col, args.label_col])
+
+model = AutoModelForTokenClassification.from_pretrained(
+    args.model_name, num_labels=len(label_list), id2label={i:l for i,l in enumerate(label_list)}, label2id={l:i for i,l in enumerate(label_list)}
+)
+
+data_collator = DataCollatorForTokenClassification(tokenizer)
+
+training_args = TrainingArguments(
+    output_dir=args.output_dir,
+    evaluation_strategy="epoch",
+    learning_rate=args.lr,
+    per_device_train_batch_size=args.batch_size,
+    per_device_eval_batch_size=args.batch_size,
+    num_train_epochs=args.epochs,
+    weight_decay=0.01,
+    load_best_model_at_end=True,
+    metric_for_best_model="f1",
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_ds,
+    eval_dataset=eval_ds,
+    tokenizer=tokenizer,
+    data_collator=data_collator,
+    compute_metrics=lambda p: compute_metrics_ner(p, label_list),
+)
+
+trainer.train()
+trainer.save_model(args.output_dir)
+tokenizer.save_pretrained(args.output_dir)