admesh/agentic-intent-classifier

1import sys2from pathlib import Path3 4import torch5from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments6 7BASE_DIR = Path(__file__).resolve().parent.parent8if str(BASE_DIR) not in sys.path:9    sys.path.insert(0, str(BASE_DIR))10 11from config import IAB_HEAD_CONFIG12from training.common import (13    build_balanced_class_weights,14    compute_classification_metrics,15    load_labeled_rows,16    prepare_dataset,17    write_json,18)19 20 21class WeightedTrainer(Trainer):22    def __init__(self, *args, class_weights: torch.Tensor | None = None, **kwargs):23        super().__init__(*args, **kwargs)24        self.class_weights = class_weights25 26    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):27        labels = inputs.pop("labels")28        outputs = model(**inputs)29        logits = outputs.get("logits")30        weight = self.class_weights.to(logits.device) if self.class_weights is not None else None31        loss_fct = torch.nn.CrossEntropyLoss(weight=weight)32        loss = loss_fct(logits.view(-1, model.config.num_labels), labels.view(-1))33        return (loss, outputs) if return_outputs else loss34 35 36train_rows = load_labeled_rows(37    IAB_HEAD_CONFIG.split_paths["train"],38    IAB_HEAD_CONFIG.label_field,39    IAB_HEAD_CONFIG.label2id,40)41val_rows = load_labeled_rows(42    IAB_HEAD_CONFIG.split_paths["val"],43    IAB_HEAD_CONFIG.label_field,44    IAB_HEAD_CONFIG.label2id,45)46test_rows = load_labeled_rows(47    IAB_HEAD_CONFIG.split_paths["test"],48    IAB_HEAD_CONFIG.label_field,49    IAB_HEAD_CONFIG.label2id,50)51 52tokenizer = AutoTokenizer.from_pretrained(IAB_HEAD_CONFIG.model_name)53 54train_dataset = prepare_dataset(train_rows, tokenizer, IAB_HEAD_CONFIG.max_length)55val_dataset = prepare_dataset(val_rows, tokenizer, IAB_HEAD_CONFIG.max_length)56test_dataset = prepare_dataset(test_rows, tokenizer, IAB_HEAD_CONFIG.max_length)57class_weights = build_balanced_class_weights(train_rows, len(IAB_HEAD_CONFIG.labels))58 59model = AutoModelForSequenceClassification.from_pretrained(60    IAB_HEAD_CONFIG.model_name,61    num_labels=len(IAB_HEAD_CONFIG.labels),62    id2label=IAB_HEAD_CONFIG.id2label,63    label2id=IAB_HEAD_CONFIG.label2id,64)65 66training_args = TrainingArguments(67    output_dir=str(IAB_HEAD_CONFIG.model_dir),68    eval_strategy="epoch",69    save_strategy="no",70    logging_strategy="epoch",71    num_train_epochs=3,72    per_device_train_batch_size=8,73    per_device_eval_batch_size=16,74    learning_rate=2e-5,75    weight_decay=0.01,76    report_to="none",77)78 79trainer = WeightedTrainer(80    model=model,81    args=training_args,82    train_dataset=train_dataset,83    eval_dataset=val_dataset,84    compute_metrics=compute_classification_metrics,85    class_weights=class_weights,86)87 88print(f"Loaded IAB splits: train={len(train_rows)} val={len(val_rows)} test={len(test_rows)}")89print(90    "IAB class weights summary:",91    {92        "min": round(float(class_weights.min().item()), 4),93        "max": round(float(class_weights.max().item()), 4),94        "mean": round(float(class_weights.mean().item()), 4),95    },96)97trainer.train()98val_metrics = trainer.evaluate(eval_dataset=val_dataset, metric_key_prefix="val")99test_metrics = trainer.evaluate(eval_dataset=test_dataset, metric_key_prefix="test")100print(val_metrics)101print(test_metrics)102 103IAB_HEAD_CONFIG.model_dir.mkdir(parents=True, exist_ok=True)104model.save_pretrained(IAB_HEAD_CONFIG.model_dir)105tokenizer.save_pretrained(IAB_HEAD_CONFIG.model_dir)106write_json(107    IAB_HEAD_CONFIG.model_dir / "train_metrics.json",108    {109        "head": IAB_HEAD_CONFIG.slug,110        "train_count": len(train_rows),111        "val_count": len(val_rows),112        "test_count": len(test_rows),113        "label_count": len(IAB_HEAD_CONFIG.labels),114        "val_metrics": val_metrics,115        "test_metrics": test_metrics,116    },117)118