Upload folder using huggingface_hub

modeling_havelock.py

@@ -5,6 +5,84 @@ import torch.nn as nn
 from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel
 
 
+class MultiLabelCRF(nn.Module):
+    """Independent CRF per marker type for multi-label BIO tagging."""
+
+    def __init__(self, num_types: int) -> None:
+        super().__init__()
+        self.num_types = num_types
+        self.transitions = nn.Parameter(torch.empty(num_types, 3, 3))
+        self.start_transitions = nn.Parameter(torch.empty(num_types, 3))
+        self.end_transitions = nn.Parameter(torch.empty(num_types, 3))
+        # Placeholder — will be overwritten by loaded weights if present
+        self.register_buffer("emission_bias", torch.zeros(1, 1, 1, 3))
+        self._reset_parameters()
+
+    def _reset_parameters(self) -> None:
+        nn.init.uniform_(self.transitions, -0.1, 0.1)
+        nn.init.uniform_(self.start_transitions, -0.1, 0.1)
+        nn.init.uniform_(self.end_transitions, -0.1, 0.1)
+        with torch.no_grad():
+            self.transitions.data[:, 0, 2] = -10000.0
+            self.start_transitions.data[:, 2] = -10000.0
+
+    def _apply_emission_bias(self, emissions: torch.Tensor) -> torch.Tensor:
+        if self.emission_bias is not None:
+            return emissions + self.emission_bias
+
+        return emissions
+
+    def decode(self, emissions: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        """Viterbi decoding.
+
+        Args:
+            emissions: (batch, seq, num_types, 3)
+            mask: (batch, seq) boolean
+
+        Returns: (batch, seq, num_types) best tag sequences
+        """
+        # Apply emission bias before decoding
+        emissions = self._apply_emission_bias(emissions)
+        batch, seq, num_types, _ = emissions.shape
+
+        # Reshape to (batch*num_types, seq, 3)
+        em = emissions.permute(0, 2, 1, 3).reshape(batch * num_types, seq, 3)
+        mk = mask.unsqueeze(1).expand(-1, num_types, -1).reshape(batch * num_types, seq)
+        BT = batch * num_types
+
+        # Expand params across batch
+        trans = (
+            self.transitions.unsqueeze(0).expand(batch, -1, -1, -1).reshape(BT, 3, 3)
+        )
+
+        start = self.start_transitions.unsqueeze(0).expand(batch, -1, -1).reshape(BT, 3)
+        end = self.end_transitions.unsqueeze(0).expand(batch, -1, -1).reshape(BT, 3)
+
+        arange = torch.arange(BT, device=em.device)
+        score = start + em[:, 0]
+        history: list[torch.Tensor] = []
+
+        for i in range(1, seq):
+            broadcast = score.unsqueeze(2) + trans + em[:, i].unsqueeze(1)
+            best_score, best_prev = broadcast.max(dim=1)
+            score = torch.where(mk[:, i].unsqueeze(1), best_score, score)
+            history.append(best_prev)
+
+        score = score + end
+        _, best_last = score.max(dim=1)
+
+        best_paths = torch.zeros(BT, seq, dtype=torch.long, device=em.device)
+        seq_lengths = mk.sum(dim=1).long()
+        best_paths[arange, seq_lengths - 1] = best_last
+
+        for i in range(seq - 2, -1, -1):
+            prev_tag = history[i][arange, best_paths[:, i + 1]]
+            should_update = i < (seq_lengths - 1)
+            best_paths[:, i] = torch.where(should_update, prev_tag, best_paths[:, i])
+
+        return best_paths.reshape(batch, num_types, seq).permute(0, 2, 1)
+
+
 class HavelockTokenConfig(PretrainedConfig):
     """Config that wraps any backbone config + our custom fields."""