Trying to count forced tokens

Author: riedgar-ms

guidance/models/_guidance_engine_metrics.py

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

guidance/models/_model.py

@@ -687,6 +687,7 @@ def next(self, logits):
                     self._sampled_token = self.tokenizer.tokens[self._sampled_token_ind]
                     self._new_bytes_prob = 1.0
                     self._was_forced = True
+                    self.metrics.forced_tokens += 1
 
                 # we are at the end of the grammar
                 elif next_byte_mask_sum == 0:
@@ -1472,6 +1473,9 @@ def _run_stateless(self, stateless_function, temperature=0.0, top_p=1.0, n=1):
         lm.engine_metrics.generated_tokens += (
             self.engine.metrics.generated_tokens - metrics_before.generated_tokens
         )
+        lm.engine_metrics.forced_tokens += (
+            self.engine.metrics.forced_tokens - metrics_before.forced_tokens
+        )
 
         logger.debug("finish Model._run_stateless")
 

tests/library/test_gen.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from guidance import gen, models
+from guidance import gen, models, select
 
 
 def test_basic():
@@ -78,17 +78,29 @@ def test_metrics_smoke(selected_model: models.Model):
     lm.reset_metrics()
 
     lm += "abc"
+    print(f"{lm.engine_metrics=}")
     lm += gen("first", max_tokens=1)
+    print(f"{lm.engine_metrics=}")
     assert lm.engine_metrics.generated_tokens == 1
 
     lm += "efg"
     lm += gen("second", max_tokens=1)
+    print(f"{lm.engine_metrics=}")
     assert lm.engine_metrics.generated_tokens == 2
 
     assert lm.current_token_count >= (
         lm.engine_metrics.prompt_tokens + lm.engine_metrics.generated_tokens
     )
 
+def test_metrics_select(selected_model: models.Model):
+    lm = selected_model
+    lm.reset_metrics()
+
+    lm += "This is a great day to "
+    lm += select(["ride a bike", "row a boat", "go for a swim"])
+    print(f"lm={str(lm)}")
+    print(f"{lm.engine_metrics=}")
+    assert False
 
 def test_unicode(selected_model):
     # black makes this test ugly -- easier to read with fmt: off