Merge branch 'main' into stateless

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.
@@ -200,20 +205,27 @@ declare module "_aici" {
      */
     constructor();
 
+    toString(): string;
+
     add(t: number): void;
     delete(t: number): void;
     has(t: number): boolean;
     clear(): void;
 
     /**
-     * Number of all tokens (not only in the set).
+     * Number of all possible tokens (regardless of whether they are in the set or not).
      */
     length: number;
 
     /**
      * Include or exclude all tokens from the set.
      */
     setAll(value: boolean): void;
+
+    /**
+     * Number of tokens in the set.
+     */
+    numSet(): number;
   }
 
   /**

controllers/jsctrl/src/jsctrl.rs

@@ -125,6 +125,11 @@ impl TokenSet {
         self.inner.len()
     }
 
+    pub fn toString(&self) -> String {
+        let trie = &mut GLOBAL_STATE.lock().unwrap().trie;
+        trie.token_set_dbg(&self.inner)
+    }
+
     pub fn add(&mut self, tok: u32) {
         self.inner.allow_token(tok);
     }
@@ -144,6 +149,10 @@ impl TokenSet {
     pub fn setAll(&mut self, val: bool) {
         self.inner.set_all(val);
     }
+
+    pub fn numSet(&self) -> usize {
+        self.inner.num_set()
+    }
 }
 
 impl Default for TokenSet {
@@ -261,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();

controllers/jsctrl/ts/aici.ts

@@ -340,6 +340,11 @@ export class ConstrainedToken extends NextToken {
       this._constraint = this.mkConstraint();
     }
     this._constraint.allowTokens(bias);
+    console.log("ALLOW:", bias.toString());
+    if (bias.numSet() === 0) {
+      console.log("Constraint doesn't allow any tokens; adding EOS")
+      return MidProcessResult.stop();
+    }
     return MidProcessResult.bias(bias);
   }
 
@@ -677,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)) {

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.
@@ -200,20 +205,27 @@ declare module "_aici" {
      */
     constructor();
 
+    toString(): string;
+
     add(t: number): void;
     delete(t: number): void;
     has(t: number): boolean;
     clear(): void;
 
     /**
-     * Number of all tokens (not only in the set).
+     * Number of all possible tokens (regardless of whether they are in the set or not).
      */
     length: number;
 
     /**
      * Include or exclude all tokens from the set.
      */
     setAll(value: boolean): void;
+
+    /**
+     * Number of tokens in the set.
+     */
+    numSet(): number;
   }
 
   /**

controllers/pyctrl/samples/test.py

@@ -3,6 +3,7 @@
 
 # asserts for microsoft/Orca-2-13b
 
+
 async def test_backtrack_one():
     await aici.FixedTokens("3+")
     l = aici.Label()
@@ -124,6 +125,15 @@ def inst(s: str) -> str:
     )
 
 
+async def test_prompt_backtrack():
+    await aici.FixedTokens("Some test prompt for the model to generate more text.")
+    l = aici.Label()
+    await aici.FixedTokens("And then some more text.")
+    await aici.gen_tokens(max_tokens=2)
+    await aici.FixedTokens("Now different text.", following=l)
+    await aici.gen_tokens(max_tokens=2)
+
+
 async def test_sample():
     # initialization code
     print("I'm going in the logs!")
@@ -161,6 +171,7 @@ async def test_eos():
     await aici.gen_tokens(regex=r' "[^"]+"', max_tokens=6, store_var="french")
     aici.check_vars({"french": ' "bonjour"'})
 
+
 async def test_joke():
     await aici.FixedTokens("Do you want a joke or a poem? A")
     answer = await aici.gen_text(options=[" joke", " poem"])

rllm/rllm-base/src/seq.rs

@@ -153,6 +153,8 @@ impl Sequence {
     ) {
         self.tokens.truncate(self.get_len() - backtrack);
         self.output_ptr = std::cmp::min(self.output_ptr, self.get_len());
+        // backtracking can remove some tokens from the initial prompt
+        self.prompt_len = std::cmp::min(self.prompt_len, self.get_len());
         if backtrack > 0 {
             self.output_pending.clear();
             self.output_pending.extend_from_slice(" ↩ ".as_bytes());