adding some debug logging for timings

modeling.py

@@ -2,6 +2,7 @@
 # This file is part of IceBERT POS model conversion.
 
 import logging
+import time
 from typing import List, Optional, Tuple
 
 import torch
@@ -379,9 +380,19 @@ class IceBertPosForTokenClassification(PreTrainedModel):
         Returns:
             List of sequences, each containing (category, [attributes]) per word
         """
+        # 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)
+        forward_time = time.perf_counter() - start_time
+        logger.debug(f"Forward pass took {forward_time:.4f} seconds")
 
-        return self._logits_to_labels(cat_logits, attr_logits, word_mask)
+        # Time the logits to labels conversion
+        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")
+
+        return result
 
     def predict_labels_from_text(
         self, sentences: List[List[str]], tokenizer, truncate: bool = False
@@ -434,11 +445,14 @@ class IceBertPosForTokenClassification(PreTrainedModel):
         # Get model predictions in (category, [attributes]) format
         predictions = self.predict_labels_from_text(sentences, tokenizer, truncate)
 
-        # Convert each sentence's predictions to IFD format
+        # Time the IFD conversion
+        start_time = time.perf_counter()
         ifd_predictions = []
         for sentence_predictions in predictions:
             ifd_labels = convert_predictions_to_ifd(sentence_predictions)  # (Ws,)
             ifd_predictions.append(ifd_labels)
+        ifd_conversion_time = time.perf_counter() - start_time
+        logger.debug(f"IFD conversion took {ifd_conversion_time:.4f} seconds")
 
         return ifd_predictions