File size: 4,144 Bytes
9d135a2 | 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 | """Iterative left-to-right beam search over the masked positions."""
from __future__ import annotations
from dataclasses import dataclass
from typing import List, Optional, Tuple
import torch
from .text_utils import clean_generated_text, split_sentences, strip_sentiment_marker
@dataclass
class Beam:
text: str
score: float
def _terminal_token_ids(tokenizer) -> set:
ids = set()
for punct in (".", "!", "?"):
ids.update(tokenizer.encode(punct, add_special_tokens=False))
return ids
@torch.inference_mode()
def fill_masks_beam_search(
model,
tokenizer,
masked_text: str,
sentence_index: int,
top_k: int = 40,
beam_size: int = 2,
max_iters: int = 30,
alpha: float = 0.7,
gamma: float = 0.05,
temperature: float = 0.8,
min_words: int = 3,
) -> Tuple[List[Beam], str]:
"""Fill masks one at a time, keeping ``beam_size`` hypotheses alive.
``alpha`` is the length-normalisation exponent, ``gamma`` penalises beams
that reuse a token already proposed this step, and a beam finishes early
once it emits terminal punctuation (after at least ``min_words`` fills).
Returns the ranked beams over the *whole* passage plus the rewritten
sentence extracted from the best beam.
"""
device = next(model.parameters()).device
mask_token_id = tokenizer.mask_token_id
pad_token_id = tokenizer.pad_token_id
if pad_token_id is None:
pad_token_id = tokenizer.eos_token_id
terminal_tokens = _terminal_token_ids(tokenizer)
encoded = tokenizer(masked_text, return_tensors="pt").to(device)
# (ids, cumulative score, tokens filled, finished, last token)
beams: List[tuple] = [(encoded.input_ids[0], 0.0, 0, False, None)]
for _ in range(max_iters):
active = [i for i, b in enumerate(beams) if (b[0] == mask_token_id).any() and not b[3]]
if not active:
break
active_beams = [beams[i] for i in active]
batch_ids = torch.stack([b[0] for b in active_beams]).to(device)
batch_scores = torch.tensor([b[1] for b in active_beams], device=device)
batch_filled = torch.tensor([b[2] for b in active_beams], device=device)
mask_positions = (batch_ids == mask_token_id).int().argmax(dim=1)
logits = model(input_ids=batch_ids).logits
mask_logits = logits[torch.arange(batch_ids.shape[0]), mask_positions] / temperature
probs = torch.softmax(mask_logits, dim=-1)
top_probs, top_tokens = torch.topk(probs, top_k)
top_probs = top_probs / top_probs.sum(dim=-1, keepdim=True)
log_p = torch.log(top_probs + 1e-10)
candidates = [b for i, b in enumerate(beams) if i not in active]
for i in range(len(active_beams)):
filled = batch_filled[i].item() + 1
length_norm = ((5 + filled) ** alpha) / ((5 + 1) ** alpha)
scores = (batch_scores[i] + log_p[i]) / length_norm
for j in range(top_k):
token_id = top_tokens[i, j].item()
# discourage every beam from picking the same continuation
penalty = sum(1 for c in candidates if c[4] == token_id)
score = scores[j].item() - gamma * penalty
new_ids = batch_ids[i].clone()
new_ids[mask_positions[i]] = token_id
finished = token_id in terminal_tokens and filled >= min_words
if finished:
new_ids[new_ids == mask_token_id] = pad_token_id
candidates.append((new_ids, score, filled, finished, token_id))
beams = sorted(candidates, key=lambda x: x[1], reverse=True)[:beam_size]
ranked = [Beam(tokenizer.decode(b[0], skip_special_tokens=True), float(b[1])) for b in beams]
best_text = ranked[0].text if ranked else ""
sentences_out = split_sentences(best_text)
if sentence_index < len(sentences_out):
best_sentence = sentences_out[sentence_index]
else:
best_sentence = sentences_out[-1] if sentences_out else best_text
return ranked, clean_generated_text(strip_sentiment_marker(best_sentence)).strip()
|