File size: 1,181 Bytes
5512228 | 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 | """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)
|