| """trainable span Head와 분리된 BIO/boundary/constrained decoding.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from itertools import product |
| from typing import Iterable, Sequence |
|
|
| import torch |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class DecodedSpan: |
| sentence_index: int |
| token_start: int |
| token_end: int |
| label: str |
| score: float |
| character_start: int | None = None |
| character_end: int | None = None |
|
|
| @property |
| def boundary(self) -> tuple[int, int, int]: |
| return self.sentence_index, self.token_start, self.token_end |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class BoundaryDecodeResult: |
| """Boundary threshold를 통과한 pair 수와 deterministic 선택 결과.""" |
|
|
| spans: tuple[DecodedSpan, ...] |
| candidate_count: int |
|
|
|
|
| class BIOFlatDecoder: |
| """BIO argmax를 flat/non-overlap span으로 복원한다.""" |
|
|
| def __init__(self, labels: Sequence[str]) -> None: |
| if not labels or labels[0] != "O": |
| raise ValueError("BIO labels must start with O") |
| self.labels = tuple(labels) |
|
|
| def decode( |
| self, |
| logits: torch.Tensor, |
| token_mask: torch.BoolTensor, |
| ) -> list[list[list[DecodedSpan]]]: |
| if logits.ndim != 4 or token_mask.shape != logits.shape[:3]: |
| raise ValueError("BIO decoder expects logits [B,S,T,C]") |
| probabilities = torch.softmax(logits, dim=-1) |
| ids = probabilities.argmax(dim=-1) |
| output: list[list[list[DecodedSpan]]] = [] |
| for batch_index in range(logits.shape[0]): |
| article = [] |
| for sentence_index in range(logits.shape[1]): |
| spans: list[DecodedSpan] = [] |
| active_label: str | None = None |
| start = 0 |
| scores: list[float] = [] |
| valid_positions = torch.nonzero( |
| token_mask[batch_index, sentence_index], as_tuple=False |
| ).flatten().tolist() |
| for offset in range(len(valid_positions) + 1): |
| token_index = ( |
| valid_positions[-1] + 1 |
| if offset == len(valid_positions) and valid_positions |
| else (0 if offset == len(valid_positions) else valid_positions[offset]) |
| ) |
| label = "O" if offset == len(valid_positions) else self.labels[int(ids[batch_index, sentence_index, token_index])] |
| prefix, _, kind = label.partition("-") |
| if active_label is not None and (prefix != "I" or kind != active_label): |
| spans.append(DecodedSpan(sentence_index, start, token_index, active_label, sum(scores) / len(scores))) |
| active_label, scores = None, [] |
| if prefix == "B" or (prefix == "I" and active_label is None): |
| active_label, start = kind, token_index |
| scores = [float(probabilities[batch_index, sentence_index, token_index, int(ids[batch_index, sentence_index, token_index])])] |
| elif prefix == "I" and kind == active_label: |
| scores.append(float(probabilities[batch_index, sentence_index, token_index, int(ids[batch_index, sentence_index, token_index])])) |
| article.append(spans) |
| output.append(article) |
| return output |
|
|
|
|
| class BoundaryProposalDecoder: |
| """독립 start/end probability에서 폭 제한 proposal을 만든다.""" |
|
|
| def __init__(self, labels: Sequence[str], threshold: float = 0.5, max_width: int = 96) -> None: |
| self.labels = tuple(labels) |
| self.threshold = threshold |
| self.max_width = max_width |
|
|
| def decode(self, start_logits: torch.Tensor, end_logits: torch.Tensor, mask: torch.BoolTensor) -> list[DecodedSpan]: |
| return list(self.decode_with_diagnostics(start_logits, end_logits, mask).spans) |
|
|
| def decode_with_diagnostics( |
| self, |
| start_logits: torch.Tensor, |
| end_logits: torch.Tensor, |
| mask: torch.BoolTensor, |
| ) -> BoundaryDecodeResult: |
| starts = torch.sigmoid(start_logits).detach() |
| ends = torch.sigmoid(end_logits).detach() |
| proposals = [] |
| for sentence in range(start_logits.shape[0]): |
| for label_index, label in enumerate(self.labels): |
| start_ids = torch.nonzero((starts[sentence, :, label_index] >= self.threshold) & mask[sentence], as_tuple=False).flatten().tolist() |
| end_ids = torch.nonzero((ends[sentence, :, label_index] >= self.threshold) & mask[sentence], as_tuple=False).flatten().tolist() |
| for start, end_inclusive in product(start_ids, end_ids): |
| if start <= end_inclusive and end_inclusive - start + 1 <= self.max_width: |
| score = float((starts[sentence, start, label_index] * ends[sentence, end_inclusive, label_index]).sqrt()) |
| proposals.append(DecodedSpan(sentence, start, end_inclusive + 1, label, score)) |
| return BoundaryDecodeResult(tuple(proposals), len(proposals)) |
|
|
|
|
| class GreedyOneToOneBoundaryDecoder: |
| """Confidence 순으로 start/end를 한 번씩만 쓰는 historical-style decoder. |
| |
| Exact duplicate는 endpoint 1:1 제약으로 제거한다. 서로 다른 endpoint를 사용하는 span의 |
| overlap/nesting은 허용하며, label·문장별 출력 cap으로 무제한 조합을 막는다. |
| """ |
|
|
| def __init__( |
| self, |
| labels: Sequence[str], |
| threshold: float = 0.5, |
| max_width: int = 64, |
| max_outputs_per_label_sentence: int = 4, |
| ) -> None: |
| if max_width <= 0 or max_outputs_per_label_sentence <= 0: |
| raise ValueError("Greedy boundary decoder limits must be positive") |
| self.labels = tuple(labels) |
| self.threshold = threshold |
| self.max_width = max_width |
| self.max_outputs_per_label_sentence = max_outputs_per_label_sentence |
|
|
| def decode( |
| self, |
| start_logits: torch.Tensor, |
| end_logits: torch.Tensor, |
| mask: torch.BoolTensor, |
| ) -> list[DecodedSpan]: |
| return list(self.decode_with_diagnostics(start_logits, end_logits, mask).spans) |
|
|
| def decode_with_diagnostics( |
| self, |
| start_logits: torch.Tensor, |
| end_logits: torch.Tensor, |
| mask: torch.BoolTensor, |
| ) -> BoundaryDecodeResult: |
| starts = torch.sigmoid(start_logits).detach() |
| ends = torch.sigmoid(end_logits).detach() |
| selected: list[DecodedSpan] = [] |
| candidate_count = 0 |
| for sentence in range(start_logits.shape[0]): |
| for label_index, label in enumerate(self.labels): |
| start_ids = torch.nonzero( |
| (starts[sentence, :, label_index] >= self.threshold) |
| & mask[sentence], |
| as_tuple=False, |
| ).flatten().tolist() |
| end_ids = torch.nonzero( |
| (ends[sentence, :, label_index] >= self.threshold) |
| & mask[sentence], |
| as_tuple=False, |
| ).flatten().tolist() |
| candidates: list[DecodedSpan] = [] |
| for start in start_ids: |
| for end_inclusive in end_ids: |
| width = end_inclusive - start + 1 |
| if 1 <= width <= self.max_width: |
| score = float( |
| ( |
| starts[sentence, start, label_index] |
| * ends[sentence, end_inclusive, label_index] |
| ).sqrt() |
| ) |
| candidates.append( |
| DecodedSpan( |
| sentence, |
| start, |
| end_inclusive + 1, |
| label, |
| score, |
| ) |
| ) |
| candidate_count += len(candidates) |
| used_starts: set[int] = set() |
| used_ends: set[int] = set() |
| for candidate in sorted( |
| candidates, |
| key=lambda item: ( |
| -item.score, |
| item.token_start, |
| item.token_end, |
| ), |
| ): |
| end_inclusive = candidate.token_end - 1 |
| if ( |
| candidate.token_start in used_starts |
| or end_inclusive in used_ends |
| ): |
| continue |
| selected.append(candidate) |
| used_starts.add(candidate.token_start) |
| used_ends.add(end_inclusive) |
| if ( |
| len(used_starts) |
| >= self.max_outputs_per_label_sentence |
| ): |
| break |
| return BoundaryDecodeResult(tuple(selected), candidate_count) |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class SemanticConstraintConfig: |
| threshold: float = 0.5 |
| max_outputs_per_sentence: int = 12 |
| allow_cross_label_exact_boundary: bool = True |
| allow_nested: bool = True |
| containment_margin: float = 0.08 |
| overlap_margin: float = 0.05 |
|
|
|
|
| class SemanticConstraintDecoder: |
| """confidence 순으로 prefix/suffix/merged proposition 중복을 제한한다.""" |
|
|
| def __init__(self, config: SemanticConstraintConfig | None = None) -> None: |
| self.config = config or SemanticConstraintConfig() |
|
|
| def decode(self, candidates: Iterable[DecodedSpan]) -> list[DecodedSpan]: |
| selected: list[DecodedSpan] = [] |
| counts: dict[int, int] = {} |
| for candidate in sorted(candidates, key=lambda item: (-item.score, item.token_start, item.token_end, item.label)): |
| if candidate.score < self.config.threshold: |
| continue |
| if counts.get(candidate.sentence_index, 0) >= self.config.max_outputs_per_sentence: |
| continue |
| if any(self._conflicts(candidate, accepted) for accepted in selected): |
| continue |
| selected.append(candidate) |
| counts[candidate.sentence_index] = counts.get(candidate.sentence_index, 0) + 1 |
| return sorted(selected, key=lambda item: (item.sentence_index, item.token_start, item.token_end, item.label)) |
|
|
| def _conflicts(self, candidate: DecodedSpan, accepted: DecodedSpan) -> bool: |
| if candidate.sentence_index != accepted.sentence_index: |
| return False |
| if candidate.boundary == accepted.boundary: |
| return not ( |
| self.config.allow_cross_label_exact_boundary and candidate.label != accepted.label |
| ) |
| overlap = max(candidate.token_start, accepted.token_start) < min(candidate.token_end, accepted.token_end) |
| if not overlap: |
| return False |
| candidate_contains = candidate.token_start <= accepted.token_start and accepted.token_end <= candidate.token_end |
| accepted_contains = accepted.token_start <= candidate.token_start and candidate.token_end <= accepted.token_end |
| if candidate_contains or accepted_contains: |
| if self.config.allow_nested and candidate.label != accepted.label: |
| return False |
| return candidate.score + self.config.containment_margin < accepted.score |
| return candidate.score <= accepted.score + self.config.overlap_margin |
|
|
|
|
| def attach_character_offsets( |
| spans: Iterable[DecodedSpan], |
| token_offsets: torch.LongTensor, |
| sentence_offsets: torch.LongTensor, |
| ) -> list[DecodedSpan]: |
| output = [] |
| for span in spans: |
| local_start = int(token_offsets[span.sentence_index, span.token_start, 0]) |
| local_end = int(token_offsets[span.sentence_index, span.token_end - 1, 1]) |
| article_start = int(sentence_offsets[span.sentence_index, 0]) + local_start |
| article_end = int(sentence_offsets[span.sentence_index, 0]) + local_end |
| output.append( |
| DecodedSpan( |
| span.sentence_index, |
| span.token_start, |
| span.token_end, |
| span.label, |
| span.score, |
| article_start, |
| article_end, |
| ) |
| ) |
| return output |
|
|
