src/lora_model.py

"""RoBERTa fine-tuned with LoRA (parameter-efficient) for intent classification."""

from peft import LoraConfig, PeftModel, TaskType, get_peft_model
from transformers import AutoModelForSequenceClassification, AutoTokenizer, PreTrainedTokenizerBase

BASE_MODEL_NAME = "roberta-base"
MAX_LENGTH = 256


def load_tokenizer() -> PreTrainedTokenizerBase:
    return AutoTokenizer.from_pretrained(BASE_MODEL_NAME)


def tokenize_batch(batch: dict, tokenizer: PreTrainedTokenizerBase) -> dict:
    return tokenizer(batch["text"], truncation=True, max_length=MAX_LENGTH)


def build_lora_roberta(
    num_labels: int,
    r: int = 32,
    lora_alpha: int = 64,
    lora_dropout: float = 0.05,
) -> PeftModel:
    """Wrap roberta-base for sequence classification with LoRA adapters on Q/K/V projections."""
    base_model = AutoModelForSequenceClassification.from_pretrained(BASE_MODEL_NAME, num_labels=num_labels)
    lora_config = LoraConfig(
        task_type=TaskType.SEQ_CLS,
        r=r,
        lora_alpha=lora_alpha,
        lora_dropout=lora_dropout,
        bias="none",
        target_modules=["query", "key", "value"],
    )
    return get_peft_model(base_model, lora_config)