Rethink the metrics

Author: riedgar-ms

guidance/models/_guidance_engine_metrics.py

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

guidance/models/_model.py

@@ -36,6 +36,8 @@
         "Failed to load guidance.cpp, falling back to Python mirror implementations..."
     )
     from .. import _cpp as cpp
+
+from ._guidance_engine_metrics import GuidanceEngineMetrics
 from .._rust.guidancerust import engine_start
 from .._utils import softmax, CaptureEvents
 from .._parser import EarleyCommitParser, Parser
@@ -52,8 +54,6 @@
 
 from .. import _serialization_pb2
 
-from ._guidance_engine_metrics import GuidanceEngineMetrics
-
 if TYPE_CHECKING:
     from ..library._block import ContextBlock
 
@@ -131,7 +131,6 @@ class EngineCallResponse:
     capture_groups: dict
     capture_group_log_probs: dict
     new_token_count: int
-    last_model_token_count: int
 
     def __init__(
         self,
@@ -141,15 +140,13 @@ 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.
@@ -208,6 +205,8 @@ def __init__(self, tokenizer, compute_log_probs=False):
         self._token_trie.match = True
         self._token_trie.match_version = 0
 
+        self.metrics = GuidanceEngineMetrics()
+
     def start(self, parser, grammar, ensure_bos_token=True):
         """Start processing parser state executed through the grammar.
 
@@ -542,7 +541,6 @@ 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,
@@ -554,6 +552,7 @@ def next(self, logits):
 
                         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)
 
@@ -740,9 +739,8 @@ def __call__(self, parser, grammar, ensure_bos_token=True):
         self.start(parser, grammar, ensure_bos_token)
 
         # TODO: remove this after the next release. This verifies that calling Rust works.
-        assert "def" == engine_start("abc", "def", 1)
+        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)
@@ -758,30 +756,20 @@ def __call__(self, parser, grammar, ensure_bos_token=True):
                     response_new_token_count,
                 ) = response_state
 
-                print(
-                    f"{response_is_generated=} {response_new_token_count=} {response_new_bytes=}"
-                )
-
                 yield EngineCallResponse(
                     new_bytes=response_new_bytes,
                     is_generated=response_is_generated,
                     new_bytes_prob=response_new_bytes_prob,
                     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, model_token_count = self.get_logits(
-                    token_ids, forced_bytes, current_temp
-                )
-                last_model_token_count = model_token_count
+                logits = self.get_logits(token_ids, forced_bytes, current_temp)
 
             if is_done:
-                assert last_model_token_count == 0, "Unyielded input tokens"
                 break
 
     def _tokenize_prefix(self, byte_string):
@@ -860,7 +848,7 @@ def _cleanup_tokens(self, token_ids, token_byte_positions):
 
         return token_ids, token_byte_positions
 
-    def get_logits(self, token_ids, forced_bytes, current_temp) -> np.ndarray:
+    def get_logits(self, token_ids, forced_bytes, current_temp):
         """A fake method designed to be overriden by subclasses."""
 
         # pretend to extend the KV cache and update the log probs
@@ -937,12 +925,6 @@ def __init__(self, engine, echo=True, **kwargs):
             0  # used to track the last event streaming call to enable throttling
         )
 
-        # Metrics for the model
-        self.engine_metrics = GuidanceEngineMetrics()
-
-    def reset_metrics(self):
-        self.engine_metrics = GuidanceEngineMetrics()
-
     @property
     def active_role_end(self):
         """The default end patterns we should use for `gen` calls.
@@ -1399,12 +1381,6 @@ def _run_stateless(self, stateless_function, temperature=0.0, top_p=1.0, n=1):
                 # if not self.engine.compute_log_probs:
                 #     chunk.new_bytes_prob = 1.0
 
-                if chunk.is_generated:
-                    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
                 chunk.new_bytes = delayed_bytes + chunk.new_bytes
@@ -1654,4 +1630,4 @@ def _check_dominated(node, parser, match_version, next_byte_mask):
             parser.pos = curr_pos
             if not child_dominate:
                 return False
-    return True
+    return True

guidance/models/transformers/_transformers.py

@@ -268,6 +268,8 @@ def get_logits(self, token_ids, forced_bytes, current_temp):
             self._cached_logits = (
                 model_out.logits[0, -1, : len(self.tokenizer.tokens)].cpu().numpy()
             )
+            self.metrics.model_input_tokens += len(new_token_ids)
+            self.metrics.model_output_tokens += 1
 
         return self._cached_logits, len(new_token_ids)