debug messages

Author: hudson-ai

guidance/models/_llama_cpp.py

@@ -215,7 +215,11 @@ def get_logits(self, token_ids: list[int], include_all_uncached_tokens: bool = F
                 logits_for_each_batch,
                 axis=0
             )
-            assert logits.shape[0] == len(token_ids) - num_cached
+            assert logits.shape[0] == len(token_ids) - num_cached, (
+                f"Expected {len(token_ids) - num_cached} logits, got {logits.shape[0]}."
+                f"\nEach batch had shape {', '.join(str(x.shape) for x in logits_for_each_batch)}."
+                f"\n{num_cached=}, {len(token_ids)=}, {n_batch=}"
+            )
             return logits
         else:
             return self._cached_logits

guidance/models/_transformers.py

@@ -635,7 +635,11 @@ def get_logits(self, token_ids: list[int], include_all_uncached_tokens: bool = F
                 logits_for_each_batch,
                 axis=0
             )
-            assert logits.shape[0] == len(token_ids) - num_cached
+            assert logits.shape[0] == len(token_ids) - num_cached, (
+                f"Expected {len(token_ids) - num_cached} logits, got {logits.shape[0]}."
+                f"\nEach batch had shape {', '.join(str(x.shape) for x in logits_for_each_batch)}."
+                f"\n{num_cached=}, {len(token_ids)=}, {len(new_token_ids)=}"
+            )
             return logits
         else:
             return self._cached_logits