File size: 4,111 Bytes
4d67ec7 | 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 | """
train_multilabel.py — Fine-tune bert-base-uncased for Multi-Label NEET classification
"""
import json
import os
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.metrics import f1_score, accuracy_score
import torch
from transformers import (
AutoTokenizer,
AutoModelForSequenceClassification,
TrainingArguments,
Trainer,
EarlyStoppingCallback,
)
from Classification.dataset_multilabel import NEETMultiLabelDataset
# ==============================
# CONFIG
# ==============================
CONFIG_PATH = "Classification/config.json"
DATA_PATH = "Classification/data/multilabel_generated_questions.json"
LABEL_MAP = "Classification/data/label_map.json"
with open(CONFIG_PATH, encoding="utf-8") as f:
cfg = json.load(f)
MODEL_NAME = cfg["model_name"]
MAX_LENGTH = cfg["max_length"]
BATCH_SIZE = cfg["batch_size"]
LR = cfg["learning_rate"]
EPOCHS = cfg["num_epochs"]
OUTPUT_DIR = "Classification/checkpoints_multilabel"
SEED = cfg["seed"]
print("Loading label map...")
with open(LABEL_MAP, "r", encoding="utf-8") as f:
label_map = json.load(f)
num_labels = len(label_map)
cfg["num_labels"] = num_labels
print(f"Loading data from {DATA_PATH}...")
if not os.path.exists(DATA_PATH):
raise FileNotFoundError(f"Missing {DATA_PATH}. Generate the data first.")
with open(DATA_PATH, encoding="utf-8") as f:
all_data = json.load(f)
# Stratification is tricky for multi-label, simple split used here
train_data, val_data = train_test_split(all_data, test_size=0.2, random_state=SEED)
os.makedirs("Classification/data", exist_ok=True)
train_path = "Classification/data/train_split_multilabel.json"
val_path = "Classification/data/val_split_multilabel.json"
with open(train_path, "w", encoding="utf-8") as f:
json.dump(train_data, f, indent=2)
with open(val_path, "w", encoding="utf-8") as f:
json.dump(val_data, f, indent=2)
print(f"Train: {len(train_data)} | Val: {len(val_data)}")
print(f"Loading tokenizer: {MODEL_NAME}")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
train_dataset = NEETMultiLabelDataset(train_path, tokenizer, num_labels, MAX_LENGTH)
val_dataset = NEETMultiLabelDataset(val_path, tokenizer, num_labels, MAX_LENGTH)
print(f"Loading {MODEL_NAME} for multi-label classification...")
model = AutoModelForSequenceClassification.from_pretrained(
MODEL_NAME,
num_labels=num_labels,
problem_type="multi_label_classification"
)
# Custom metric function
def compute_metrics(eval_pred):
logits, labels = eval_pred
# Apply sigmoid
sigmoid = torch.nn.Sigmoid()
probs = sigmoid(torch.Tensor(logits))
# Threshold at 0.5
preds = np.zeros(probs.shape)
preds[np.where(probs >= 0.5)] = 1
# Macro F1 is standard for multi-label
macro_f1 = f1_score(labels, preds, average='macro', zero_division=0)
acc = accuracy_score(labels, preds)
return {"accuracy": acc, "f1_macro": macro_f1}
training_args = TrainingArguments(
output_dir=OUTPUT_DIR,
num_train_epochs=EPOCHS,
per_device_train_batch_size=BATCH_SIZE,
per_device_eval_batch_size=BATCH_SIZE,
learning_rate=LR,
warmup_ratio=cfg["warmup_ratio"],
weight_decay=cfg["weight_decay"],
evaluation_strategy=cfg["evaluation_strategy"],
save_strategy=cfg["save_strategy"],
load_best_model_at_end=cfg["load_best_model_at_end"],
metric_for_best_model="f1_macro",
logging_steps=cfg["logging_steps"],
seed=SEED,
report_to="none",
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=val_dataset,
tokenizer=tokenizer,
compute_metrics=compute_metrics,
callbacks=[EarlyStoppingCallback(early_stopping_patience=2)],
)
print("\nStarting multi-label training...")
trainer.train()
best_model_dir = os.path.join(OUTPUT_DIR, "best_model")
trainer.save_model(best_model_dir)
tokenizer.save_pretrained(best_model_dir)
print(f"\n✅ Best multi-label model saved to: {best_model_dir}")
metrics = trainer.evaluate()
print(f"\nFinal Eval Metrics: {metrics}")
|