Warm up every bucket at load so the first request per size is not ~10x slower

src/gliner_decide_coreml/router.py

@@ -9,6 +9,7 @@ from dataclasses import dataclass
 from pathlib import Path
 
 import coremltools as ct
+import numpy as np
 
 from .decoding import decode, head_probabilities, merge_chunks
 from .encoding import Encoded, encode, load_collator, task_labels, to_bucket
@@ -57,6 +58,17 @@ class DecideRouter:
             for length in BUCKETS
         }
         self.chunk_overlap_words = chunk_overlap_words
+        self._warm_up()
+
+    def _warm_up(self):
+        """Run each function once: the first prediction per function is ~10x slower (device setup)."""
+        for length, model in self.models.items():
+            model.predict({
+                "input_ids": np.full((1, length), self.pad_id, dtype=np.int32),
+                "attention_mask": np.ones((1, length), dtype=np.int32),
+                "marker_indices": np.zeros((1, MAX_HEADS, MAX_OPTIONS), dtype=np.int32),
+                "marker_mask": np.ones((1, MAX_HEADS, MAX_OPTIONS), dtype=np.float32),
+            })
 
     def count_tokens(self, text: str, tasks: dict) -> int:
         return encode(self.collator, text, tasks).length