26 lines
871 B
Python
26 lines
871 B
Python
from __future__ import annotations
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from models.fact_prototype import CanonicalFactPrototype
|
|
|
|
|
|
def canonical_prototype_loss(
|
|
states: torch.Tensor,
|
|
targets: list[torch.Tensor],
|
|
assignments: list[tuple[torch.Tensor, torch.Tensor]],
|
|
prototype_encoder: CanonicalFactPrototype,
|
|
) -> torch.Tensor:
|
|
vectors = []
|
|
labels = []
|
|
for batch_index, (slot_index, target_index) in enumerate(assignments):
|
|
if len(slot_index):
|
|
vectors.append(states[batch_index, slot_index])
|
|
labels.append(targets[batch_index].to(states.device)[target_index])
|
|
if not vectors:
|
|
return states.sum() * 0.0
|
|
state_tensor = F.normalize(torch.cat(vectors), dim=-1)
|
|
prototype_tensor = prototype_encoder(torch.cat(labels))
|
|
return (1.0 - (state_tensor * prototype_tensor).sum(-1)).mean()
|