File size: 6,928 Bytes
54c3e65 f11438f 54c3e65 f11438f 54c3e65 f11438f 54c3e65 f11438f 54c3e65 f11438f 54c3e65 f11438f 54c3e65 | 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 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 | from __future__ import annotations
from pathlib import Path
from typing import Any
import torch
from transformers import AutoModelForMaskedLM, AutoTokenizer
from .domain import CandidateScoringError, RerankRequest
MODEL_ID = "ku-nlp/deberta-v2-tiny-japanese"
MODEL_REVISION = "0427645e8cf44ee83b4a0b5f4498274d89e02adb"
class DebertaCandidateScorer:
"""Score only the candidate span with masked-LM pseudo-log-likelihood."""
def __init__(
self,
*,
model_id: str = MODEL_ID,
revision: str = MODEL_REVISION,
device: str = "cpu",
cache_dir: Path | None = None,
local_files_only: bool = False,
max_length: int = 256,
max_mask_batch: int = 16,
) -> None:
self._device = torch.device(device)
self._model_id = model_id
self._load_options: dict[str, Any] = {
"revision": revision,
"local_files_only": local_files_only,
}
if cache_dir is not None:
self._load_options["cache_dir"] = str(cache_dir)
self._tokenizer: Any | None = None
self._model: Any | None = None
self._requested_max_length = max_length
self._max_length = max_length
self._max_mask_batch = max_mask_batch
def load(self) -> None:
"""Load and pin model assets before latency measurement."""
self._ensure_loaded()
@torch.inference_mode()
def score_candidates(self, request: RerankRequest) -> list[float]:
tokenizer, model = self._loaded_components()
input_rows: list[torch.Tensor] = []
attention_rows: list[torch.Tensor] = []
positions: list[int] = []
target_ids: list[int] = []
candidate_indexes: list[int] = []
for candidate_index, candidate in enumerate(request.candidates):
encoding, target_index = self._encode_candidate(
tokenizer,
request.left_context,
request.right_context,
candidate.surface,
)
word_ids = encoding.word_ids(batch_index=0)
target_positions = [
position for position, word_id in enumerate(word_ids) if word_id == target_index
]
if not target_positions:
raise CandidateScoringError(
"candidate_outside_window",
"candidate was outside the model token window",
)
for position in target_positions:
masked_ids = encoding["input_ids"][0].clone()
target_id = int(masked_ids[position])
if target_id == tokenizer.unk_token_id:
raise CandidateScoringError(
"unknown_token",
"candidate contains an unknown model token",
)
if target_id in tokenizer.all_special_ids:
raise CandidateScoringError(
"special_token",
"candidate resolves to a special model token",
)
masked_ids[position] = tokenizer.mask_token_id
input_rows.append(masked_ids)
attention_rows.append(encoding["attention_mask"][0])
positions.append(position)
target_ids.append(target_id)
candidate_indexes.append(candidate_index)
sums = [0.0] * len(request.candidates)
counts = [0] * len(request.candidates)
for start in range(0, len(input_rows), self._max_mask_batch):
stop = start + self._max_mask_batch
ids = torch.nn.utils.rnn.pad_sequence(
input_rows[start:stop],
batch_first=True,
padding_value=tokenizer.pad_token_id,
).to(self._device)
attention = torch.nn.utils.rnn.pad_sequence(
attention_rows[start:stop], batch_first=True, padding_value=0
).to(self._device)
logits = model(input_ids=ids, attention_mask=attention).logits
row_count = logits.shape[0]
row_indexes = torch.arange(row_count, device=self._device)
chunk_positions = torch.tensor(
positions[start:stop], dtype=torch.long, device=self._device
)
chunk_targets = torch.tensor(
target_ids[start:stop], dtype=torch.long, device=self._device
)
target_logits = logits[row_indexes, chunk_positions]
log_probs = torch.log_softmax(target_logits, dim=-1)
values = log_probs[row_indexes, chunk_targets].cpu().tolist()
for offset, value in enumerate(values):
candidate_index = candidate_indexes[start + offset]
sums[candidate_index] += float(value)
counts[candidate_index] += 1
if any(count == 0 for count in counts):
raise CandidateScoringError(
"candidate_unscored",
"at least one candidate could not be scored",
)
return [total / count for total, count in zip(sums, counts, strict=True)]
def _ensure_loaded(self) -> None:
if self._tokenizer is not None and self._model is not None:
return
tokenizer = AutoTokenizer.from_pretrained(self._model_id, **self._load_options)
model = AutoModelForMaskedLM.from_pretrained(self._model_id, **self._load_options)
model.to(self._device)
model.eval()
model_limit = int(getattr(model.config, "max_position_embeddings", 512))
self._tokenizer = tokenizer
self._model = model
self._max_length = min(self._requested_max_length, model_limit)
def _loaded_components(self) -> tuple[Any, Any]:
self._ensure_loaded()
assert self._tokenizer is not None
assert self._model is not None
return self._tokenizer, self._model
def _encode_candidate(
self,
tokenizer: Any,
left_context: tuple[str, ...],
right_context: tuple[str, ...],
surface: str,
) -> tuple[Any, int]:
left = list(left_context)
right = list(right_context)
while True:
words = [*left, surface, *right]
target_index = len(left)
encoding = tokenizer(
words,
is_split_into_words=True,
return_tensors="pt",
)
if encoding["input_ids"].shape[1] <= self._max_length:
return encoding, target_index
if not left and not right:
raise CandidateScoringError(
"candidate_too_long",
"candidate is longer than the model token window",
)
if len(left) >= len(right) and left:
left.pop(0)
elif right:
right.pop()
|