Latest attempt to get consistent token results

Author: riedgar-ms

guidance/models/_guidance_engine_metrics.py

@@ -4,3 +4,4 @@
 class GuidanceEngineMetrics(BaseModel):
     generated_tokens: NonNegativeInt = 0
     forced_tokens: NonNegativeInt = 0
+    model_input_tokens: NonNegativeInt = 0

guidance/models/_model.py

@@ -131,6 +131,7 @@ class EngineCallResponse:
     capture_groups: dict
     capture_group_log_probs: dict
     new_token_count: int
+    last_model_token_count: int
 
     def __init__(
         self,
@@ -140,13 +141,15 @@ def __init__(
         capture_groups,
         capture_group_log_probs,
         new_token_count,
+        last_model_token_count,
     ):
         self.new_bytes = new_bytes
         self.is_generated = is_generated
         self.new_bytes_prob = new_bytes_prob
         self.capture_groups = capture_groups
         self.capture_group_log_probs = capture_group_log_probs
         self.new_token_count = new_token_count
+        self.last_model_token_count = last_model_token_count
 
     def _to_proto(self):
         """Converts an EngineCallResponse object to its Protobuf representation.
@@ -739,6 +742,7 @@ def __call__(self, parser, grammar, ensure_bos_token=True):
         # TODO: remove this after the next release. This verifies that calling Rust works.
         assert "def" == engine_start("abc", "def", 1)
 
+        last_model_token_count = 0
         logits = None
         while True:
             is_done, logits_state, response_state = self.next(logits)
@@ -765,13 +769,19 @@ def __call__(self, parser, grammar, ensure_bos_token=True):
                     capture_groups=response_capture_groups,
                     capture_group_log_probs=response_capture_group_log_probs,
                     new_token_count=response_new_token_count,
+                    last_model_token_count=last_model_token_count,
                 )
+                last_model_token_count = 0
 
             if logits_state is not None:
                 token_ids, forced_bytes, current_temp = logits_state
-                logits = self.get_logits(token_ids, forced_bytes, current_temp)
+                logits, model_token_count = self.get_logits(
+                    token_ids, forced_bytes, current_temp
+                )
+                last_model_token_count = model_token_count
 
             if is_done:
+                assert last_model_token_count == 0, "Unyielded input tokens"
                 break
 
     def _tokenize_prefix(self, byte_string):
@@ -1393,6 +1403,7 @@ def _run_stateless(self, stateless_function, temperature=0.0, top_p=1.0, n=1):
                     self.engine_metrics.generated_tokens += chunk.new_token_count
                 else:
                     self.engine_metrics.forced_tokens += chunk.new_token_count
+                self.engine_metrics.model_input_tokens += chunk.last_model_token_count
 
                 # convert the bytes to a string (delaying if we don't yet have a valid unicode string)
                 lm.token_count += chunk.new_token_count

guidance/models/transformers/_transformers.py

@@ -269,7 +269,7 @@ def get_logits(self, token_ids, forced_bytes, current_temp):
                 model_out.logits[0, -1, : len(self.tokenizer.tokens)].cpu().numpy()
             )
 
-        return self._cached_logits
+        return self._cached_logits, len(new_token_ids)
 
 
 class Transformers(Model):

tests/library/test_gen.py

@@ -138,8 +138,11 @@ def test_metrics_alt_expressions(selected_model: models.Model):
     assert str(lm) == str(lm2)
     assert lm.engine_metrics.generated_tokens == 10
     assert lm2.engine_metrics.generated_tokens == 10
-    assert lm.engine_metrics.forced_tokens == 0
-    assert lm2.engine_metrics.forced_tokens == 0
+
+    assert (
+        lm.engine_metrics.forced_tokens + lm.engine_metrics.model_input_tokens
+        == lm2.engine_metrics.forced_tokens + lm2.engine_metrics.model_input_tokens
+    )
 
 
 def test_unicode(selected_model):