From ca5fd2b7d2e803c0d8c7191eef41f41c99725416 Mon Sep 17 00:00:00 2001
From: Michal Moskal <michal@moskal.me>
Date: Thu, 22 Feb 2024 01:21:41 +0000
Subject: [PATCH] usability fixes

---
 controllers/aici_abi/src/toktree.rs | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/controllers/aici_abi/src/toktree.rs b/controllers/aici_abi/src/toktree.rs
index c64055a2..72150d12 100644
--- a/controllers/aici_abi/src/toktree.rs
+++ b/controllers/aici_abi/src/toktree.rs
@@ -205,7 +205,14 @@ impl TokTrie {
     }
 
     pub fn token_dbg(&self, idx: u32) -> String {
-        format!("{:?}[{}]", self.token_str(idx), idx)
+        if idx == self.info.tok_eos {
+            "EOS".to_string()
+        } else if idx as usize >= self.vocab_size() {
+            format!("OOB[{}]", idx)
+        } else {
+            // format!("{:?}[{}]", self.token_str(idx), idx)
+            format!("{:?}", self.token_str(idx))
+        }
     }
 
     pub fn token_str(&self, idx: u32) -> String {
@@ -471,6 +478,8 @@ impl TokTrie {
             }
         }
         r.trie_finished();
+        // revert the fake token
+        toks.disallow_token(defl_tok);
     }
 }