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 (12    DECISION_PHASE_DIFFICULTY_DATA_DIR,13    DECISION_PHASE_HEAD_CONFIG,14    DECISION_PHASE_TRAINING_WEIGHTS,15    FULL_INTENT_TAXONOMY_DATA_DIR,16)17from training.common import (18    build_label_weight_tensor,19    compute_classification_metrics,20    load_labeled_rows,21    load_labeled_rows_from_paths,22    prepare_dataset,23    write_json,24)25 26 27class WeightedTrainer(Trainer):28    def __init__(self, *args, class_weights: torch.Tensor | None = None, **kwargs):29        super().__init__(*args, **kwargs)30        self.class_weights = class_weights31 32    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):33        labels = inputs.pop("labels")34        outputs = model(**inputs)35        logits = outputs.get("logits")36        weight = self.class_weights.to(logits.device) if self.class_weights is not None else None37        loss_fct = torch.nn.CrossEntropyLoss(weight=weight)38        loss = loss_fct(logits.view(-1, model.config.num_labels), labels.view(-1))39        return (loss, outputs) if return_outputs else loss40 41train_rows = load_labeled_rows_from_paths(42    [43        DECISION_PHASE_HEAD_CONFIG.split_paths["train"],44        FULL_INTENT_TAXONOMY_DATA_DIR / "train.jsonl",45        DECISION_PHASE_DIFFICULTY_DATA_DIR / "train.jsonl",46    ],47    DECISION_PHASE_HEAD_CONFIG.label_field,48    DECISION_PHASE_HEAD_CONFIG.label2id,49)50val_rows = load_labeled_rows_from_paths(51    [52        DECISION_PHASE_HEAD_CONFIG.split_paths["val"],53        FULL_INTENT_TAXONOMY_DATA_DIR / "val.jsonl",54        DECISION_PHASE_DIFFICULTY_DATA_DIR / "val.jsonl",55    ],56    DECISION_PHASE_HEAD_CONFIG.label_field,57    DECISION_PHASE_HEAD_CONFIG.label2id,58)59test_rows = load_labeled_rows(60    DECISION_PHASE_HEAD_CONFIG.split_paths["test"],61    DECISION_PHASE_HEAD_CONFIG.label_field,62    DECISION_PHASE_HEAD_CONFIG.label2id,63)64 65tokenizer = AutoTokenizer.from_pretrained(DECISION_PHASE_HEAD_CONFIG.model_name)66 67train_dataset = prepare_dataset(train_rows, tokenizer, DECISION_PHASE_HEAD_CONFIG.max_length)68val_dataset = prepare_dataset(val_rows, tokenizer, DECISION_PHASE_HEAD_CONFIG.max_length)69test_dataset = prepare_dataset(test_rows, tokenizer, DECISION_PHASE_HEAD_CONFIG.max_length)70class_weights = build_label_weight_tensor(DECISION_PHASE_HEAD_CONFIG.labels, DECISION_PHASE_TRAINING_WEIGHTS)71 72model = AutoModelForSequenceClassification.from_pretrained(73    DECISION_PHASE_HEAD_CONFIG.model_name,74    num_labels=len(DECISION_PHASE_HEAD_CONFIG.labels),75    id2label=DECISION_PHASE_HEAD_CONFIG.id2label,76    label2id=DECISION_PHASE_HEAD_CONFIG.label2id,77)78 79training_args = TrainingArguments(80    output_dir=str(DECISION_PHASE_HEAD_CONFIG.model_dir),81    eval_strategy="epoch",82    save_strategy="no",83    logging_strategy="epoch",84    num_train_epochs=4,85    per_device_train_batch_size=4,86    per_device_eval_batch_size=4,87    learning_rate=2e-5,88    weight_decay=0.01,89    report_to="none",90)91 92trainer = WeightedTrainer(93    model=model,94    args=training_args,95    train_dataset=train_dataset,96    eval_dataset=val_dataset,97    compute_metrics=compute_classification_metrics,98    class_weights=class_weights,99)100 101print(102    f"Loaded decision_phase splits: train={len(train_rows)} val={len(val_rows)} test={len(test_rows)}"103)104print(f"Decision phase weights: {[round(float(x), 3) for x in class_weights.tolist()]}")105trainer.train()106val_metrics = trainer.evaluate(eval_dataset=val_dataset, metric_key_prefix="val")107test_metrics = trainer.evaluate(eval_dataset=test_dataset, metric_key_prefix="test")108print(val_metrics)109print(test_metrics)110 111DECISION_PHASE_HEAD_CONFIG.model_dir.mkdir(parents=True, exist_ok=True)112model.save_pretrained(DECISION_PHASE_HEAD_CONFIG.model_dir)113tokenizer.save_pretrained(DECISION_PHASE_HEAD_CONFIG.model_dir)114write_json(115    DECISION_PHASE_HEAD_CONFIG.model_dir / "train_metrics.json",116    {117        "head": DECISION_PHASE_HEAD_CONFIG.slug,118        "train_count": len(train_rows),119        "val_count": len(val_rows),120        "test_count": len(test_rows),121        "val_metrics": val_metrics,122        "test_metrics": test_metrics,123    },124)125