File size: 1,201 Bytes
6fd3740 | 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 | """Span pointer decoding for the v3 copy head."""
from __future__ import annotations
import torch
def best_span(start: torch.Tensor, end: torch.Tensor, max_len: int = 32):
"""Highest p(start) * p(end) with start <= end < start + max_len -> (s, e, prob)."""
ps, pe = start.float().softmax(-1), end.float().softmax(-1)
Q = ps.shape[0]
i = torch.arange(Q, device=ps.device)
ok = (i[None, :] >= i[:, None]) & (i[None, :] < i[:, None] + max_len)
score = (ps[:, None] * pe[None, :]).masked_fill(~ok, 0)
flat = int(score.argmax())
return flat // Q, flat % Q, float(score.flatten()[flat])
def _cls(ch: str) -> str:
return "a" if ch.isalpha() else "d" if ch.isdigit() else ""
def snap_word(text: str, s: int, e: int) -> tuple[int, int]:
"""Widen the char span [s, e) so it never cuts a word or number in half ("US Doll" -> "US Dollars").
Only letter-letter and digit-digit edges are extended, so "100" stays "100" inside "100km"."""
while s > 0 and s < e and _cls(text[s]) and _cls(text[s - 1]) == _cls(text[s]):
s -= 1
while e < len(text) and e > s and _cls(text[e - 1]) and _cls(text[e]) == _cls(text[e - 1]):
e += 1
return s, e
|