File size: 7,015 Bytes
9df97a2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 | #!/usr/bin/env python3
"""Fine-tune a recruiter-domain NER model from annotated CV data.
Expected input formats:
- JSONL with tokens + ner_tags
- JSON with records containing `tokens` and `ner_tags`
This script is intentionally defensive: if `transformers` is unavailable,
it still validates and exports a prepared dataset snapshot for later training.
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
from typing import Any, Dict, List, Sequence, Tuple
try:
from datasets import Dataset # type: ignore
from transformers import AutoModelForTokenClassification, AutoTokenizer, Trainer, TrainingArguments # type: ignore
TRANSFORMERS_AVAILABLE = True
except Exception:
Dataset = None
AutoModelForTokenClassification = None
AutoTokenizer = None
Trainer = None
TrainingArguments = None
TRANSFORMERS_AVAILABLE = False
LABEL_NAMES = ["O", "B-PER", "I-PER", "B-ORG", "I-ORG", "B-SKILL", "I-SKILL", "B-EDU", "I-EDU"]
DEFAULT_MODEL = "dslim/bert-base-NER"
def load_annotations(path: Path) -> List[Dict[str, Any]]:
"""Load token/label annotations from JSON or JSONL."""
if not path.exists():
raise FileNotFoundError(f"Annotation file not found: {path}")
text = path.read_text(encoding="utf-8").strip()
if not text:
return []
if path.suffix.lower() == ".jsonl":
records = []
for line in text.splitlines():
if line.strip():
records.append(json.loads(line))
return records
payload = json.loads(text)
if isinstance(payload, list):
return payload
if isinstance(payload, dict) and isinstance(payload.get("data"), list):
return payload["data"]
raise ValueError("Unsupported annotation format")
def validate_records(records: Sequence[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""Keep only records with aligned tokens/labels."""
valid_records: List[Dict[str, Any]] = []
for record in records:
tokens = record.get("tokens")
ner_tags = record.get("ner_tags")
if not isinstance(tokens, list) or not isinstance(ner_tags, list):
continue
if len(tokens) != len(ner_tags) or not tokens:
continue
valid_records.append({"tokens": [str(token) for token in tokens], "ner_tags": [str(tag) for tag in ner_tags]})
return valid_records
def build_label_map(records: Sequence[Dict[str, Any]]) -> Tuple[Dict[str, int], Dict[int, str]]:
labels = {"O"}
for record in records:
labels.update(str(tag) for tag in record.get("ner_tags", []))
ordered_labels = ["O"] + sorted(label for label in labels if label != "O")
label2id = {label: index for index, label in enumerate(ordered_labels)}
id2label = {index: label for label, index in label2id.items()}
return label2id, id2label
def tokenize_and_align_labels(examples: Dict[str, List[Any]], tokenizer, label2id: Dict[str, int]) -> Dict[str, Any]:
tokenized_inputs = tokenizer(
examples["tokens"],
is_split_into_words=True,
truncation=True,
padding=False,
)
aligned_labels = []
for batch_index, labels in enumerate(examples["ner_tags"]):
word_ids = tokenized_inputs.word_ids(batch_index=batch_index)
previous_word_id = None
label_ids: List[int] = []
for word_id in word_ids:
if word_id is None:
label_ids.append(-100)
elif word_id != previous_word_id:
label_ids.append(label2id.get(labels[word_id], label2id["O"]))
else:
label = labels[word_id]
if label.startswith("B-"):
label = "I-" + label[2:]
label_ids.append(label2id.get(label, label2id["O"]))
previous_word_id = word_id
aligned_labels.append(label_ids)
tokenized_inputs["labels"] = aligned_labels
return tokenized_inputs
def export_prepared_dataset(records: Sequence[Dict[str, Any]], output_dir: Path) -> Path:
output_dir.mkdir(parents=True, exist_ok=True)
export_path = output_dir / "recruiter_ner_prepared_dataset.json"
export_path.write_text(json.dumps(list(records), indent=2, ensure_ascii=False), encoding="utf-8")
return export_path
def train_model(records: Sequence[Dict[str, Any]], model_name: str, output_dir: Path, epochs: int) -> Path:
if not TRANSFORMERS_AVAILABLE:
raise RuntimeError("transformers/datasets are not available in this environment")
label2id, id2label = build_label_map(records)
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForTokenClassification.from_pretrained(
model_name,
num_labels=len(label2id),
label2id=label2id,
id2label=id2label,
)
dataset = Dataset.from_list(list(records))
tokenized_dataset = dataset.map(lambda batch: tokenize_and_align_labels(batch, tokenizer, label2id), batched=True)
output_dir.mkdir(parents=True, exist_ok=True)
training_args = TrainingArguments(
output_dir=str(output_dir),
learning_rate=2e-5,
per_device_train_batch_size=8,
per_device_eval_batch_size=8,
num_train_epochs=epochs,
weight_decay=0.01,
logging_steps=10,
save_strategy="epoch",
report_to=[],
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_dataset,
tokenizer=tokenizer,
)
trainer.train()
trainer.save_model(str(output_dir))
tokenizer.save_pretrained(str(output_dir))
return output_dir
def main() -> int:
parser = argparse.ArgumentParser(description="Fine-tune recruiter NER model")
parser.add_argument("--annotations", required=True, help="Path to JSON/JSONL annotation file")
parser.add_argument("--output-dir", default="models/recruiter_ner", help="Training output directory")
parser.add_argument("--model-name", default=DEFAULT_MODEL, help="Base HF model to fine-tune")
parser.add_argument("--epochs", type=int, default=3, help="Training epochs")
parser.add_argument("--export-only", action="store_true", help="Prepare dataset without training")
args = parser.parse_args()
annotation_path = Path(args.annotations).expanduser().resolve()
output_dir = Path(args.output_dir).expanduser().resolve()
records = validate_records(load_annotations(annotation_path))
if not records:
raise SystemExit("No valid annotation records found")
export_path = export_prepared_dataset(records, output_dir)
print(f"Prepared dataset exported to {export_path}")
if args.export_only or not TRANSFORMERS_AVAILABLE:
print("Training skipped: exporting only or transformers unavailable")
return 0
trained_dir = train_model(records, args.model_name, output_dir, args.epochs)
print(f"Model trained and saved to {trained_dir}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|