add tokenRepr in jsctrl

microsoft · Mar 15, 2024 · 6262f05 · 6262f05
1 parent 977e50a
commit 6262f05
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 0 deletions.
diff --git a/controllers/jsctrl/samples/aici-types.d.ts b/controllers/jsctrl/samples/aici-types.d.ts
@@ -143,6 +143,11 @@ declare module "_aici" {
    */
   function detokenize(tokens: number[]): Buffer;
 
+  /**
+   * Return debug string representation of a given token index
+   */
+  function tokenRepr(token: number): string;
+
   /**
    * Return identifier of the current sequence.
    * Most useful with fork_group parameter in mid_process() callback.

diff --git a/controllers/jsctrl/src/jsctrl.rs b/controllers/jsctrl/src/jsctrl.rs
@@ -270,6 +270,12 @@ mod aici_mod {
         Buffer(bytes)
     }
 
+    #[rquickjs::function]
+    pub fn tokenRepr(token: TokenId) -> String {
+        let trie = &mut GLOBAL_STATE.lock().unwrap().trie;
+        trie.token_dbg(token)
+    }
+
     #[rquickjs::function]
     pub fn getVar(name: String) -> Option<Buffer> {
         let name = name.as_str();

diff --git a/controllers/jsctrl/ts/aici.ts b/controllers/jsctrl/ts/aici.ts
@@ -682,6 +682,8 @@ export async function genTokens(options: GenOptions): Promise<Token[]> {
     const tokens = await next_token.run();
     res.push(...tokens);
 
+    console.log("GEN-STEP:", tokens.map(t => _aici.tokenRepr(t)).join(", "));
+
     const text = detokenize(res).decode();
 
     if (stopAt !== undefined && text.includes(stopAt)) {

diff --git a/controllers/jsctrl/ts/native.d.ts b/controllers/jsctrl/ts/native.d.ts
@@ -143,6 +143,11 @@ declare module "_aici" {
    */
   function detokenize(tokens: number[]): Buffer;
 
+  /**
+   * Return debug string representation of a given token index
+   */
+  function tokenRepr(token: number): string;
+
   /**
    * Return identifier of the current sequence.
    * Most useful with fork_group parameter in mid_process() callback.