| from __future__ import annotations |
|
|
| import torch |
| from torch import Tensor, nn |
|
|
|
|
| class RawExpert(nn.Module): |
| def __init__(self, hidden_size: int, decision_size: int): |
| super().__init__() |
| self.net = nn.Sequential( |
| nn.Linear(hidden_size * 2, decision_size), |
| nn.GELU(), |
| nn.LayerNorm(decision_size), |
| nn.Linear(decision_size, decision_size), |
| ) |
|
|
| def forward(self, state: Tensor, question: Tensor) -> Tensor: |
| return self.net(torch.cat((state, question), dim=-1)) |
|
|
|
|
| class CandidateScorer(nn.Module): |
| """Shared pointwise scorer; candidate permutation only permutes its outputs.""" |
|
|
| def __init__(self, decision_size: int): |
| super().__init__() |
| self.net = nn.Sequential( |
| nn.Linear(decision_size * 4, decision_size), |
| nn.GELU(), |
| nn.Linear(decision_size, 1), |
| ) |
|
|
| def forward(self, decision: Tensor, candidates: Tensor) -> Tensor: |
| if decision.ndim == 1: |
| decision = decision.unsqueeze(0) |
| if candidates.ndim == 2: |
| candidates = candidates.unsqueeze(0).expand(decision.size(0), -1, -1) |
| decision = decision.unsqueeze(1).expand(-1, candidates.size(1), -1) |
| features = torch.cat( |
| (decision, candidates, decision * candidates, torch.abs(decision - candidates)), dim=-1 |
| ) |
| return self.net(features).squeeze(-1).squeeze(0) |
|
|
|
|
| class NoulHead(nn.Module): |
| def __init__(self, decision_size: int): |
| super().__init__() |
| self.linear = nn.Linear(decision_size, 1) |
|
|
| def forward(self, decision: Tensor) -> Tensor: |
| return self.linear(decision).squeeze(-1) |
|
|
