From bb5840871de0c8ed701b3c044117450fc9efd67c Mon Sep 17 00:00:00 2001
From: Michal Moskal <michal@moskal.me>
Date: Fri, 2 Feb 2024 00:49:26 +0000
Subject: [PATCH] fix tokenset size in jsctrl; fixes #47

---
 jsctrl/src/jsctrl.rs | 8 ++++----
 rllm-cuda/server.sh  | 5 +++++
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/jsctrl/src/jsctrl.rs b/jsctrl/src/jsctrl.rs
index 1da6818a..5e6cf98c 100644
--- a/jsctrl/src/jsctrl.rs
+++ b/jsctrl/src/jsctrl.rs
@@ -114,10 +114,10 @@ impl<'js> Trace<'js> for TokenSet {
 impl TokenSet {
     #[qjs(constructor)]
     pub fn new() -> TokenSet {
-        // TODO set the size
-        let mut inner = SimpleVob::new();
-        inner.resize(32000);
-        TokenSet { inner }
+        let trie = &mut GLOBAL_STATE.lock().unwrap().trie;
+        TokenSet {
+            inner: trie.alloc_token_set(),
+        }
     }
 
     #[qjs(get)]
diff --git a/rllm-cuda/server.sh b/rllm-cuda/server.sh
index 900eb52c..0e33d49f 100755
--- a/rllm-cuda/server.sh
+++ b/rllm-cuda/server.sh
@@ -39,6 +39,11 @@ if [ "$1" = "--trace" ] ; then
     shift
 fi
 
+if [ "$1" = "--trace-rt" ] ; then
+    R_LOG=info,rllm=trace,aicirt=trace
+    shift
+fi
+
 if [ "$1" = "--loop" ] ; then
     LOOP=1
     shift