| """Classifier with LoRA and class-weighted loss.""" |
|
|
| import torch |
| import torch.nn as nn |
| from transformers import AutoModelForSequenceClassification, Trainer |
| from peft import LoraConfig, get_peft_model, TaskType |
|
|
| from config import CONFIG |
|
|
|
|
| def load_model(): |
| """Load base model and wrap with LoRA adapter for efficient fine-tuning.""" |
| base_model = AutoModelForSequenceClassification.from_pretrained( |
| CONFIG["model_name"], |
| num_labels=CONFIG["num_labels"], |
| ) |
|
|
| if CONFIG.get("gradient_checkpointing", False): |
| base_model.gradient_checkpointing_enable() |
|
|
| lora_config = LoraConfig( |
| task_type=TaskType.SEQ_CLS, |
| r=CONFIG.get("lora_r", 16), |
| lora_alpha=CONFIG.get("lora_alpha", 32), |
| lora_dropout=CONFIG.get("lora_dropout", 0.05), |
| target_modules=CONFIG.get("lora_target_modules", ["query", "value"]), |
| modules_to_save=["classifier"], |
| ) |
|
|
| model = get_peft_model(base_model, lora_config) |
| model.print_trainable_parameters() |
|
|
| return model |
|
|
|
|
| class WeightedTrainer(Trainer): |
| """Trainer subclass that uses class-weighted CrossEntropyLoss. |
| |
| Compatible with transformers v4.x and v5.x. |
| """ |
|
|
| def __init__(self, class_weights=None, **kwargs): |
| super().__init__(**kwargs) |
| if class_weights is not None: |
| self._class_weights = torch.tensor(class_weights, dtype=torch.float32) |
| else: |
| self._class_weights = None |
|
|
| def compute_loss(self, model, inputs, return_outputs=False, **kwargs): |
| labels = inputs.get("labels") |
| outputs = model(**inputs) |
| logits = outputs.get("logits") |
|
|
| if self._class_weights is not None: |
| weight = self._class_weights.to(logits.device) |
| loss_fn = nn.CrossEntropyLoss(weight=weight) |
| else: |
| loss_fn = nn.CrossEntropyLoss() |
|
|
| loss = loss_fn(logits.view(-1, self.model.config.num_labels), labels.view(-1)) |
| return (loss, outputs) if return_outputs else loss |
|
|
