| """Predicted dependency and constituent structure over induced spans.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| import math |
|
|
| import torch |
| from torch import nn |
|
|
| from strata.modeling.ph_pat.config import PHPATConfig |
| from strata.modeling.ph_pat.span_compiler import SpanCompilerOutput |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class DependencyChartOutput: |
| nodes: torch.Tensor |
| head_logits: torch.Tensor |
| relation_logits: torch.Tensor |
| constituent_logits: torch.Tensor |
| segment_ids: torch.Tensor |
| valid_mask: torch.Tensor |
|
|
|
|
| class DependencyChart(nn.Module): |
| def __init__(self, config: PHPATConfig) -> None: |
| super().__init__() |
| self.config = config |
| self.head_query = nn.Linear(config.d_model, config.d_model, bias=False) |
| self.head_key = nn.Linear(config.d_model, config.d_model, bias=False) |
| self.relation = nn.Linear(2 * config.d_model, config.dependency_relations) |
| self.constituent = nn.Linear(config.d_model, config.constituent_categories) |
| self.node_update = nn.Sequential( |
| nn.Linear(config.d_model, config.d_model), |
| nn.SiLU(), |
| nn.Linear(config.d_model, config.d_model), |
| ) |
|
|
| def forward(self, spans: SpanCompilerOutput) -> DependencyChartOutput: |
| nodes = spans.nodes + self.node_update(spans.nodes) |
| query = self.head_query(nodes) |
| key = self.head_key(nodes) |
| logits = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.config.d_model) |
| same_segment = spans.segment_ids.unsqueeze(-1) == spans.segment_ids.unsqueeze(-2) |
| valid_pair = spans.valid_mask.unsqueeze(-1) & spans.valid_mask.unsqueeze(-2) & same_segment |
| logits = logits.masked_fill(~valid_pair, torch.finfo(logits.dtype).min) |
|
|
| head_weights = torch.softmax(logits.float(), dim=-1).to(nodes.dtype) |
| head_weights = torch.nan_to_num(head_weights) |
| parent = torch.matmul(head_weights, nodes) |
| relation_logits = self.relation(torch.cat((nodes, parent), dim=-1)) |
| return DependencyChartOutput( |
| nodes=nodes, |
| head_logits=logits, |
| relation_logits=relation_logits, |
| constituent_logits=self.constituent(nodes), |
| segment_ids=spans.segment_ids, |
| valid_mask=spans.valid_mask, |
| ) |
|
|
|
|
| __all__ = ["DependencyChart", "DependencyChartOutput"] |
|
|
