File size: 4,357 Bytes
98188cc | 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 | """Mask builders for targeted masked fine-tuning."""
from __future__ import annotations
from typing import Iterable, Sequence
import re
import numpy as np
import torch
import torch.nn.functional as F
DEFAULT_NER_LABELS = {"PERSON", "ORG", "GPE", "LOC", "EMAIL", "PHONE", "DATE"}
EMAIL_RE = re.compile(r"[a-zA-Z0-9_.+-]+@[a-zA-Z0-9-]+(?:\.[a-zA-Z0-9-]+)+")
PHONE_RE = re.compile(r"(?<!\d)(?:\+?1[-.\s]?)?(?:\(?\d{3}\)?[-.\s]?)\d{3}[-.\s]\d{4}(?!\d)")
def char_span_to_token_indices(
offset_mapping: Sequence[tuple[int, int]],
char_start: int,
char_end: int,
) -> list[int]:
"""Return token indices whose character offsets overlap a character span."""
token_indices: list[int] = []
for idx, (token_start, token_end) in enumerate(offset_mapping):
if token_start == token_end:
continue
if token_start < char_end and token_end > char_start:
token_indices.append(idx)
return token_indices
def ner_mask(
offset_mapping: Sequence[tuple[int, int]],
text: str,
nlp_model,
ner_labels: Iterable[str] | None = None,
) -> torch.Tensor:
"""Build a boolean token mask from spaCy NER spans."""
labels = set(ner_labels or DEFAULT_NER_LABELS)
mask = torch.zeros(len(offset_mapping), dtype=torch.bool)
doc = nlp_model(text)
for ent in doc.ents:
if ent.label_ not in labels:
continue
token_indices = char_span_to_token_indices(
offset_mapping,
ent.start_char,
ent.end_char,
)
if token_indices:
mask[token_indices] = True
for pattern, label in ((EMAIL_RE, "EMAIL"), (PHONE_RE, "PHONE")):
if label not in labels:
continue
for match in pattern.finditer(text):
token_indices = char_span_to_token_indices(offset_mapping, match.start(), match.end())
if token_indices:
mask[token_indices] = True
return mask
def random_mask(
attention_mask: torch.Tensor,
probability: float,
generator: torch.Generator | None = None,
) -> torch.Tensor:
"""Mask a random subset of non-padding tokens."""
if probability <= 0:
return torch.zeros_like(attention_mask, dtype=torch.bool)
probs = torch.rand(
attention_mask.shape,
generator=generator,
device=attention_mask.device,
)
return (probs < probability) & attention_mask.bool()
def _token_nll(model, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
outputs = model(input_ids=input_ids, attention_mask=attention_mask)
logits = outputs.logits
shift_logits = logits[:, :-1, :].contiguous()
shift_labels = input_ids[:, 1:].contiguous()
losses = F.cross_entropy(
shift_logits.view(-1, shift_logits.size(-1)),
shift_labels.view(-1),
reduction="none",
).view(input_ids.size(0), -1)
padded = torch.zeros_like(input_ids, dtype=losses.dtype)
padded[:, 1:] = losses
return padded
def mia_mask(
input_ids: torch.Tensor,
attention_mask: torch.Tensor,
target_model,
reference_model,
threshold_percentile: float = 75,
min_score: float = 0.0,
) -> torch.Tensor:
"""Mask tokens the target predicts better than the frozen base model."""
device = next(target_model.parameters()).device
input_ids = input_ids.to(device)
attention_mask = attention_mask.to(device)
reference_model = reference_model.to(device)
target_was_training = target_model.training
target_model.eval()
reference_model.eval()
with torch.no_grad():
target_loss = _token_nll(target_model, input_ids, attention_mask)
ref_loss = _token_nll(reference_model, input_ids, attention_mask)
score = ref_loss - target_loss
valid = attention_mask.bool()
valid[:, 0] = False
valid_scores = score[valid]
if valid_scores.numel() == 0:
if target_was_training:
target_model.train()
return torch.zeros_like(attention_mask, dtype=torch.bool)
threshold = max(
float(min_score),
float(np.percentile(valid_scores.detach().float().cpu().numpy(), threshold_percentile)),
)
mask = score > threshold
if target_was_training:
target_model.train()
return (mask & valid).cpu()
|