From f49791ce1289ee302344073014f3a17967306446 Mon Sep 17 00:00:00 2001
From: Michal Moskal <michal@moskal.me>
Date: Thu, 14 Nov 2024 15:06:34 -0800
Subject: [PATCH] track more parser stats

---
 parser/src/earley/parser.rs        | 21 +++++++++++++++++++--
 parser/src/tokenparser.rs          | 22 ++++++++++++++++++----
 sample_parser/src/sample_parser.rs | 10 ++++++++--
 3 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/parser/src/earley/parser.rs b/parser/src/earley/parser.rs
index 1c7de11..265f3b4 100644
--- a/parser/src/earley/parser.rs
+++ b/parser/src/earley/parser.rs
@@ -15,6 +15,7 @@ use std::{
 
 use anyhow::{bail, ensure, Result};
 use derivre::{RegexAst, StateID};
+use instant::Instant;
 use serde::{Deserialize, Serialize};
 use toktrie::{Recognizer, SimpleVob, SpecialToken, TokEnv, TokTrie, TokenId};
 
@@ -61,8 +62,8 @@ pub struct ParserStats {
     pub num_lex_errors: usize,
     pub num_lexemes: usize,
     pub all_items: usize,
-
     pub lexer_cost: u64,
+    pub compute_time_us: u64,
 }
 
 impl ParserStats {
@@ -74,8 +75,21 @@ impl ParserStats {
             num_lexemes: self.num_lexemes - previous.num_lexemes,
             num_lex_errors: self.num_lex_errors - previous.num_lex_errors,
             all_items: self.all_items - previous.all_items,
-
             lexer_cost: self.lexer_cost - previous.lexer_cost,
+            compute_time_us: self.compute_time_us - previous.compute_time_us,
+        }
+    }
+
+    pub fn max(&self, other: &ParserStats) -> ParserStats {
+        ParserStats {
+            rows: self.rows.max(other.rows),
+            definitive_bytes: self.definitive_bytes.max(other.definitive_bytes),
+            lexer_ops: self.lexer_ops.max(other.lexer_ops),
+            num_lexemes: self.num_lexemes.max(other.num_lexemes),
+            num_lex_errors: self.num_lex_errors.max(other.num_lex_errors),
+            all_items: self.all_items.max(other.all_items),
+            lexer_cost: self.lexer_cost.max(other.lexer_cost),
+            compute_time_us: self.compute_time_us.max(other.compute_time_us),
         }
     }
 }
@@ -407,6 +421,7 @@ impl ParserState {
         computer: &dyn BiasComputer,
         start: &[u8],
     ) -> SimpleVob {
+        let t0 = Instant::now();
         let dfa = &mut shared.lexer.dfa;
         dfa.set_fuel(self.limits.step_lexer_fuel);
         dfa.set_max_states(self.limits.max_lexer_states);
@@ -438,6 +453,8 @@ impl ParserState {
             set.allow_token(computer.trie().eos_token());
         }
 
+        self.stats.compute_time_us += t0.elapsed().as_micros() as u64;
+
         set
     }
 
diff --git a/parser/src/tokenparser.rs b/parser/src/tokenparser.rs
index 225956a..f8e413a 100644
--- a/parser/src/tokenparser.rs
+++ b/parser/src/tokenparser.rs
@@ -22,6 +22,8 @@ pub struct TokenParser {
     pub logger: Logger,
     pub limits: ParserLimits,
     pub bias_computer: Arc<dyn BiasComputer>,
+    last_step_stats: ParserStats,
+    max_step_stats: ParserStats,
     pending_bogus_backtrack: u32,
     // sampling any of these will pop the parser stack:
     pop_tokens: Option<SimpleVob>,
@@ -93,6 +95,8 @@ impl TokenParser {
             inference_caps,
             limits,
             pending_bogus_backtrack: 0,
+            max_step_stats: ParserStats::default(),
+            last_step_stats: ParserStats::default(),
             mid_process_start_time,
             mid_process_was_accepting: false,
             no_bias_this_mid_process: false,
@@ -133,6 +137,14 @@ impl TokenParser {
         self.parser.stats()
     }
 
+    pub fn last_step_stats(&self) -> &ParserStats {
+        &self.last_step_stats
+    }
+
+    pub fn max_step_stats(&self) -> &ParserStats {
+        &self.max_step_stats
+    }
+
     pub fn num_tokens(&self) -> usize {
         self.llm_tokens.len()
     }
@@ -633,7 +645,6 @@ impl TokenParser {
             return StepResult::noop();
         }
 
-        let pre = instant::Instant::now();
         let pre_stats = self.parser.stats().clone();
         let mut set = self
             .parser
@@ -643,7 +654,9 @@ impl TokenParser {
             let err = format!("lexer error: {}", err);
             return self.stop(&err, StopReason::LexerTooComplex);
         }
-        self.last_bias_time = pre.elapsed();
+        self.last_bias_time = Duration::from_micros(p_stats.compute_time_us);
+        self.last_step_stats = p_stats.clone();
+        self.max_step_stats = self.max_step_stats.max(&p_stats);
 
         if inner_accepting {
             let mut all_accepting = true;
@@ -682,9 +695,10 @@ impl TokenParser {
 
         infoln!(
             self,
-            "step-stats: {:?}; {} lex fuel; {}",
-            start_time.elapsed(),
+            "step-stats: {}us; {} lex fuel; {} items; {}",
+            start_time.elapsed().as_micros(),
             p_stats.lexer_cost,
+            p_stats.all_items,
             self.parser.lexer_stats(),
         );
 
diff --git a/sample_parser/src/sample_parser.rs b/sample_parser/src/sample_parser.rs
index b7e46f6..0550ffb 100644
--- a/sample_parser/src/sample_parser.rs
+++ b/sample_parser/src/sample_parser.rs
@@ -83,11 +83,15 @@ fn main() {
             black_box(constraint.temperature);
             let sampled_token = tokens[idx];
 
+            let p_stats = constraint.parser.last_step_stats();
             println!(
-                "SAMPLE {}: {} {}",
+                "SAMPLE {}: {} {}; stats: {} lex, {} rows, {} us",
                 idx,
                 sampled_token,
-                tok_env.tok_trie().token_dbg(sampled_token)
+                tok_env.tok_trie().token_dbg(sampled_token),
+                p_stats.lexer_cost,
+                p_stats.all_items,
+                p_stats.compute_time_us,
             );
             Some(sampled_token)
         } else {
@@ -130,6 +134,8 @@ fn main() {
     send_output(&constraint.flush_logs());
     // the stop reason should be likely also sent to the user
     println!("Stop reason: {:?}", constraint.parser.stop_reason());
+
+    println!("Max step stats: {:?}", constraint.parser.max_step_stats());
 }
 
 fn read_file_to_string(filename: &str) -> String {