Tweak where stats are grabbed

Author: riedgar-ms

guidance/models/_guidance_engine_metrics.py

@@ -2,6 +2,5 @@
 
 
 class GuidanceEngineMetrics(BaseModel):
-    prompt_tokens: NonNegativeInt = 0
     generated_tokens: NonNegativeInt = 0
     forced_tokens: NonNegativeInt = 0

guidance/models/_model.py

@@ -539,18 +539,18 @@ def next(self, logits):
                         # self._captured_log_prob_data.update(new_captured_log_prob_data)
                         # yield out, self._is_generated, self._new_bytes_prob, self._captured_data, self._captured_log_prob_data, self._token_count - self._last_token_count # note that we don't capture groups until a complete parse right now...
 
+                        self._token_count += 1  # note we only update this for tokens that emit non-hidden content
                         response_state = (
                             out,
                             is_generated,
                             self._new_bytes_prob if self.compute_log_probs else 1.0,
                             self._captured_data,
                             self._captured_log_prob_data,
-                            self._token_count - self._last_token_count,
+                            self._token_count - self._last_token_count + 1,
                         )
 
                         self._last_token_count = self._token_count
                         self._hidden_count = 0
-                        self._token_count += 1  # note we only update this for tokens that emit non-hidden content
                     else:
                         self._hidden_count -= len(new_bytes)
 

tests/library/test_gen.py

@@ -101,11 +101,16 @@ def test_metrics_select(selected_model: models.Model):
     lm += select(["ride a bike", "row a boat", "go for a swim"])
     print(f"lm={str(lm)}")
     print(f"{lm.engine_metrics=}")
+    assert lm.engine_metrics.forced_tokens > 0
+    assert lm.engine_metrics.generated_tokens > 0
+    assert lm.engine_metrics.forced_tokens > lm.engine_metrics.generated_tokens
+    prev_stats = lm.engine_metrics.copy()
     lm += " and afterwards "
     lm += select(["walk to town", "walk to a show"])
     print(f"lm={str(lm)}")
     print(f"{lm.engine_metrics=}")
-    assert False
+    assert lm.engine_metrics.forced_tokens > prev_stats.forced_tokens
+    assert lm.engine_metrics.generated_tokens > prev_stats.generated_tokens
 
 
 def test_unicode(selected_model):