Format modeling

modeling.py

@@ -36,7 +36,9 @@ class MultiLabelTokenClassificationHead(nn.Module):
         # (*, H) -> (*, C)
         self.cat_proj = nn.Linear(self.hidden_size, self.num_categories)
         # (*, H + C) -> (*, A)
-        self.out_proj = nn.Linear(self.hidden_size + self.num_categories, self.num_labels)
+        self.out_proj = nn.Linear(
+            self.hidden_size + self.num_categories, self.num_labels
+        )
 
     def forward(self, features: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
         """
@@ -119,9 +121,15 @@ class IceBertPosForTokenClassification(PreTrainedModel):
 
         # Create group attribute indices tensor: (G x max_group_size)
         # Instead of dict lookups, we can index directly: group_attr_indices[group_id, :]
-        max_group_size = max(len(labels) for labels in schema.group_name_to_labels.values())
-        self.group_attr_indices = torch.full((num_groups, max_group_size), -1, dtype=torch.long, device=device)
-        self.group_sizes = torch.zeros(num_groups, dtype=torch.long, device=device)  # (G,)
+        max_group_size = max(
+            len(labels) for labels in schema.group_name_to_labels.values()
+        )
+        self.group_attr_indices = torch.full(
+            (num_groups, max_group_size), -1, dtype=torch.long, device=device
+        )
+        self.group_sizes = torch.zeros(
+            num_groups, dtype=torch.long, device=device
+        )  # (G,)
 
         for group_idx, group_name in enumerate(schema.group_names):
             group_labels = schema.group_name_to_labels[group_name]
@@ -179,7 +187,9 @@ class IceBertPosForTokenClassification(PreTrainedModel):
             cat_logits: Category logits (B x W x C)
             attr_logits: Attribute logits (B x W x A)
         """
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
 
         # Get RoBERTa outputs
         outputs = self.roberta(
@@ -197,7 +207,9 @@ class IceBertPosForTokenClassification(PreTrainedModel):
         hidden_states = outputs[0]  # (B x L x H)
 
         # (B x L x H) -> (Wt x H)
-        word_embeddings = self._aggregate_subword_tokens(hidden_states, word_mask, attention_mask)
+        word_embeddings = self._aggregate_subword_tokens(
+            hidden_states, word_mask, attention_mask
+        )
 
         # (Wt x H) -> (Wt x C), (Wt x A)
         cat_logits, attr_logits = self.classifier(word_embeddings)
@@ -209,7 +221,10 @@ class IceBertPosForTokenClassification(PreTrainedModel):
         return cat_logits, attr_logits
 
     def _aggregate_subword_tokens(
-        self, sequence_output: torch.Tensor, word_mask: torch.Tensor, attention_mask: torch.Tensor
+        self,
+        sequence_output: torch.Tensor,
+        word_mask: torch.Tensor,
+        attention_mask: torch.Tensor,
     ) -> torch.Tensor:
         """
         Average subword tokens within each word to get word-level representations.
@@ -242,13 +257,17 @@ class IceBertPosForTokenClassification(PreTrainedModel):
 
             # Get word starts for this sequence
             seq_word_mask = word_mask[b, valid_mask]  # (Lv,) - only valid positions
-            word_starts = seq_word_mask.nonzero(as_tuple=True)[0]  # (Ws,) - positions where words start
+            word_starts = seq_word_mask.nonzero(as_tuple=True)[
+                0
+            ]  # (Ws,) - positions where words start
 
             if len(word_starts) == 0:
                 continue
 
             # Assign each token to its word within this sequence
-            seq_word_indices = torch.full((len(seq_word_mask),), -1, dtype=torch.long, device=device)
+            seq_word_indices = torch.full(
+                (len(seq_word_mask),), -1, dtype=torch.long, device=device
+            )
 
             for i, start_pos in enumerate(word_starts):
                 # Find end position (next word start or end of sequence)
@@ -268,7 +287,10 @@ class IceBertPosForTokenClassification(PreTrainedModel):
         # This allows us to use scatter operations across the entire batch
         max_words_per_seq = word_mask.sum(dim=-1)  # (B,) - words per sequence
         word_offset = torch.cat(
-            [torch.zeros(1, device=device, dtype=torch.long), max_words_per_seq.cumsum(dim=0)[:-1]]
+            [
+                torch.zeros(1, device=device, dtype=torch.long),
+                max_words_per_seq.cumsum(dim=0)[:-1],
+            ]
         )  # (B,) - cumulative word offsets
 
         # Add batch offsets to make global unique indices
@@ -283,7 +305,9 @@ class IceBertPosForTokenClassification(PreTrainedModel):
         # Only use tokens that belong to words (not padding and not before first word)
         valid_word_tokens = (flat_attention.bool()) & (flat_word_indices >= 0)  # (B*L,)
         valid_output = flat_output[valid_word_tokens]  # (valid_word_tokens x H)
-        valid_word_indices = flat_word_indices[valid_word_tokens]  # (valid_word_tokens,)
+        valid_word_indices = flat_word_indices[
+            valid_word_tokens
+        ]  # (valid_word_tokens,)
 
         total_words = max_words_per_seq.sum()
         if total_words == 0:
@@ -292,11 +316,17 @@ class IceBertPosForTokenClassification(PreTrainedModel):
         # Vectorized aggregation using scatter operations
         # Sum all token embeddings that belong to the same word
         word_sums = torch.zeros(total_words, hidden_size, device=device)  # (Wt x H)
-        word_sums.scatter_add_(0, valid_word_indices.unsqueeze(1).expand(-1, hidden_size), valid_output)
+        word_sums.scatter_add_(
+            0, valid_word_indices.unsqueeze(1).expand(-1, hidden_size), valid_output
+        )
 
         # Count how many tokens belong to each word (for averaging)
         word_counts = torch.zeros(total_words, device=device)  # (Wt,)
-        word_counts.scatter_add_(0, valid_word_indices, torch.ones_like(valid_word_indices, dtype=torch.float))
+        word_counts.scatter_add_(
+            0,
+            valid_word_indices,
+            torch.ones_like(valid_word_indices, dtype=torch.float),
+        )
 
         # Compute average: word_embedding = sum_of_tokens / count_of_tokens
         word_counts = torch.clamp(word_counts, min=1.0)  # Prevent division by zero
@@ -304,7 +334,9 @@ class IceBertPosForTokenClassification(PreTrainedModel):
 
         return word_features
 
-    def _reshape_to_batch_format(self, logits: torch.Tensor, nwords: torch.Tensor) -> torch.Tensor:
+    def _reshape_to_batch_format(
+        self, logits: torch.Tensor, nwords: torch.Tensor
+    ) -> torch.Tensor:
         """
         Reshape concatenated word predictions back to padded batch format.
 
@@ -357,7 +389,9 @@ class IceBertPosForTokenClassification(PreTrainedModel):
 
         # Debug logging to match fairseq model
         logger.debug(f"Encoded tokens: {input_ids}")  # (L,)
-        logger.debug(f"Decoded tokens: {tokenizer.convert_ids_to_tokens(input_ids.tolist())}")
+        logger.debug(
+            f"Decoded tokens: {tokenizer.convert_ids_to_tokens(input_ids.tolist())}"
+        )
         logger.debug(f"Word IDs: {word_ids}")  # (L,)
         logger.debug(f"Word mask: {word_mask}")
 
@@ -365,7 +399,10 @@ class IceBertPosForTokenClassification(PreTrainedModel):
 
     @torch.no_grad()
     def predict_labels(
-        self, input_ids: torch.Tensor, attention_mask: torch.Tensor, word_mask: torch.Tensor
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        word_mask: torch.Tensor,
     ) -> List[List[Tuple[str, List[str]]]]:
         """
         Predict POS labels for input sequences.
@@ -382,7 +419,9 @@ class IceBertPosForTokenClassification(PreTrainedModel):
         """
         # Time the forward pass
         start_time = time.perf_counter()
-        cat_logits, attr_logits = self.forward(input_ids=input_ids, attention_mask=attention_mask, word_mask=word_mask)
+        cat_logits, attr_logits = self.forward(
+            input_ids=input_ids, attention_mask=attention_mask, word_mask=word_mask
+        )
         forward_time = time.perf_counter() - start_time
         logger.debug(f"Forward pass took {forward_time:.4f} seconds")
 
@@ -390,7 +429,9 @@ class IceBertPosForTokenClassification(PreTrainedModel):
         start_time = time.perf_counter()
         result = self._logits_to_labels(cat_logits, attr_logits, word_mask)
         logits_to_labels_time = time.perf_counter() - start_time
-        logger.debug(f"Logits to labels conversion took {logits_to_labels_time:.4f} seconds")
+        logger.debug(
+            f"Logits to labels conversion took {logits_to_labels_time:.4f} seconds"
+        )
 
         return result
 
@@ -414,19 +455,31 @@ class IceBertPosForTokenClassification(PreTrainedModel):
         all_word_masks = []
 
         for words in sentences:
-            input_ids, attention_mask, word_mask = self.prepare_inputs(words, tokenizer, truncate)
+            input_ids, attention_mask, word_mask = self.prepare_inputs(
+                words, tokenizer, truncate
+            )
             all_input_ids.append(input_ids)
             all_attention_masks.append(attention_mask)
             all_word_masks.append(word_mask)
 
         # Pad sequences to same length
-        batch_input_ids = pad_sequence(all_input_ids, batch_first=True, padding_value=tokenizer.pad_token_id)
-        batch_attention_mask = pad_sequence(all_attention_masks, batch_first=True, padding_value=0)
-        batch_word_mask = pad_sequence(all_word_masks, batch_first=True, padding_value=0)
+        batch_input_ids = pad_sequence(
+            all_input_ids, batch_first=True, padding_value=tokenizer.pad_token_id
+        )
+        batch_attention_mask = pad_sequence(
+            all_attention_masks, batch_first=True, padding_value=0
+        )
+        batch_word_mask = pad_sequence(
+            all_word_masks, batch_first=True, padding_value=0
+        )
 
-        return self.predict_labels(batch_input_ids, batch_attention_mask, batch_word_mask)
+        return self.predict_labels(
+            batch_input_ids, batch_attention_mask, batch_word_mask
+        )
 
-    def convert_labels_to_ifd(self, predictions: List[List[Tuple[str, List[str]]]]) -> List[List[str]]:
+    def convert_labels_to_ifd(
+        self, predictions: List[List[Tuple[str, List[str]]]]
+    ) -> List[List[str]]:
         """
         Convert model predictions to IFD format labels.
 
@@ -495,7 +548,10 @@ class IceBertPosForTokenClassification(PreTrainedModel):
         return word_mask
 
     def _logits_to_labels(
-        self, cat_logits: torch.Tensor, attr_logits: torch.Tensor, word_mask: torch.Tensor
+        self,
+        cat_logits: torch.Tensor,
+        attr_logits: torch.Tensor,
+        word_mask: torch.Tensor,
     ) -> List[List[Tuple[str, List[str]]]]:
         """
         Convert logits to human-readable labels using vectorized operations.
@@ -520,7 +576,9 @@ class IceBertPosForTokenClassification(PreTrainedModel):
             if nwords[b] > 0:
                 batch_word_mask[b, : nwords[b]] = True
 
-        valid_positions = batch_word_mask.flatten().nonzero(as_tuple=True)[0]  # (total_words,)
+        valid_positions = batch_word_mask.flatten().nonzero(as_tuple=True)[
+            0
+        ]  # (total_words,)
         total_words = len(valid_positions)
 
         if total_words == 0:
@@ -607,7 +665,9 @@ class IceBertPosForTokenClassification(PreTrainedModel):
                     attributes = []
                 elif cat_name == "sl" and "act" in attributes:
                     # Number and tense are not shown for sl act in mim format
-                    attributes = [attr for attr in attributes if attr not in ["1", "sing", "pres"]]
+                    attributes = [
+                        attr for attr in attributes if attr not in ["1", "sing", "pres"]
+                    ]
 
                 seq_predictions.append((cat_name, attributes))
                 word_counter += 1