File size: 813 Bytes
945de56 | 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 | """DistilBERT for sentence-pair boundary classification."""
from transformers import (
AutoTokenizer,
DistilBertForSequenceClassification,
PreTrainedTokenizerFast,
)
from src.datasets.combined_pairs_dataset import NUM_LABELS, ID2LABEL, LABEL2ID
BASE_MODEL = "distilbert-base-uncased"
def load_distilbert(
pretrained: str = BASE_MODEL,
) -> DistilBertForSequenceClassification:
"""Instantiate DistilBERT for 3-class sentence-pair classification."""
return DistilBertForSequenceClassification.from_pretrained(
pretrained,
num_labels=NUM_LABELS,
id2label=ID2LABEL,
label2id=LABEL2ID,
)
def load_distilbert_tokenizer(
pretrained: str = BASE_MODEL,
) -> PreTrainedTokenizerFast:
return AutoTokenizer.from_pretrained(pretrained, use_fast=True)
|