incremental streaming of decoder graph in playground

Author: lbeurerkellner

src/lmql/runtime/dclib/dclib_seq.py

@@ -11,6 +11,11 @@
 
 detokenize_seqs = True
 
+@dataclass
+class DecoderGraphSnapshot:
+    node_hashes = None
+    step_not_updated_count = None
+
 class DecoderGraph:
     def __init__(self):
         self.nodes = {}
@@ -19,6 +24,8 @@ def __init__(self):
 
         self.ctr = 0
 
+        self.json_snapshot: DecoderGraphSnapshot = DecoderGraphSnapshot()
+
     def add_node(self, node):
         uid = f"n{self.ctr}"
         self.ctr += 1
@@ -34,13 +41,39 @@ def set_pool(self, node, pool_name):
     def add_edge(self, from_node, to_node):
         self.edges.append((self.node_ids[from_node], self.node_ids[to_node]))
 
-    async def json(self):
+    async def json(self, diff: bool = False):
         nodes = []
+        
+        if self.json_snapshot.node_hashes is None:
+            self.json_snapshot.node_hashes = {}
+        if self.json_snapshot.step_not_updated_count is None:
+            self.json_snapshot.step_not_updated_count = {}
+        
         for k, v in self.nodes.items():
+            hash = None
+            if diff and k in self.json_snapshot.node_hashes:
+                # this will ignore changes to nodes that have not been updated for 3 steps (may miss changes to nodes that 
+                # are not updated for a long time, careful for now)
+                if k in self.json_snapshot.step_not_updated_count and self.json_snapshot.step_not_updated_count[k] > 2:
+                    continue
+                hash = await v.json_hash()
+                if self.json_snapshot.node_hashes[k] == str(hash):
+                    self.json_snapshot.step_not_updated_count[k] = self.json_snapshot.step_not_updated_count.get(k, 0) + 1
+                    continue
+                    
+                # else:
+                #     print("node changed", k, self.json_snapshot.node_hashes[k], hash)
+            if hash is None:
+                hash = await v.json_hash()
+
             nodes.append({
                 "id": k,
                 **await v.json()
             })
+            
+            if diff:
+                self.json_snapshot.node_hashes[k] = str(hash)
+                self.json_snapshot.step_not_updated_count[k] = 0
 
         return {
             "nodes": nodes,
@@ -180,6 +213,12 @@ async def text_provider():
 
         return self.tokenizer_cache[name]
 
+    async def json_hash(self):
+        o = await self.json()
+        if type(o["user_data"]) is dict and "openai-continuations" in o["user_data"]:
+            o["user_data"].pop("openai-continuations")
+        return hash(str(o))
+
     async def json(self):
         seqtext = await self.detokenized("seqtext")
         text = await self.detokenized("text")

src/lmql/runtime/dclib/lmql_adapter.py

@@ -166,7 +166,7 @@ async def query(self, prompt, mask_logits_processor, head_input_id_rewriter, act
 
         async def debug_out(decoder_step):
             if _DCLibDebugPrinter.printer is not None and dc.DecoderSequence.graph is not None:
-                data = await dc.DecoderSequence.graph.json()
+                data = await dc.DecoderSequence.graph.json(diff=True)
                 data = replace_inf_nan_with_str(data)
                 _DCLibDebugPrinter.printer.add_decoder_state(data)
             dcmodel.report_stats(_DCLibDebugPrinter.printer, decoder_step)
@@ -184,17 +184,18 @@ async def debug_out(decoder_step):
                 if step_budget is not None and decoder_step >= step_budget:
                     print("warning: step budget exceeded")
                     break
-                
-                if "performance_stats" in decoder_args:
-                    Stats.print_all()
 
                 if interrupt.check():
                     interrupt.clear()
                     raise InterruptedError("lmql.runtime.interrupt")
 
-                average_step_time = time.time() - start if average_step_time is None else (average_step_time * 0.9 + (time.time() - start) * 0.1)
-                # if decoder_step % 10 == 0:
-                #     print("step", decoder_step, "time", average_step_time)
+                average_step_time = (time.time() - start) if average_step_time is None else (average_step_time * 0.9 + (time.time() - start) * 0.1)
+
+                if "performance_stats" in decoder_args:
+                    if decoder_step % 10 == 0:
+                        Stats.print_all()
+                        print("step", decoder_step, "time", average_step_time)
+
                 start = time.time()
 
         except dc.FinishException as fe:
@@ -219,7 +220,7 @@ async def debug_out(decoder_step):
             return [self.make_decoder_head(i,n,s) for i,s in enumerate(result_sequences)]
 
     def validate_args(self, decoder_args, decoder_fct):
-        INTERNAL_ARGS = ["decoder", "dcmodel", "dclib_additional_logits_processor", "input_id_rewriter", "output_writer", "chatty_openai", "distribution_batch_size", "openai_chunksize", "step_budget", "stats"]
+        INTERNAL_ARGS = ["decoder", "dcmodel", "dclib_additional_logits_processor", "input_id_rewriter", "output_writer", "chatty_openai", "distribution_batch_size", "openai_chunksize", "step_budget", "stats", "performance_stats"]
 
         # get all arg names and kwarg names of decoder function
         decoder_arg_names = inspect.getfullargspec(decoder_fct).args

src/lmql/runtime/tokenizer.py

@@ -103,8 +103,16 @@ def decode(self, input_ids):
         if n-1 in self.detokenizer_cache.keys():
             key = str(input_ids[:-1])
             if key in self.detokenizer_cache[n-1].keys():
+                global reverse_special_token_mappings
                 # print("secondary cache hit")
-                return self.detokenizer_cache[n-1][key] + self.tokenizer_impl.decode([input_ids[-1]])
+                if input_ids[-1] >= self.tokenizer_impl.vocab_size:
+                    extended = self.detokenizer_cache[n-1][key] + "<" + reverse_special_token_mappings[input_ids[-1]] + "/>"
+                else:
+                    extended = self.detokenizer_cache[n-1][key] + self.tokenizer_impl.decode([input_ids[-1]])
+                if not n in self.detokenizer_cache.keys():
+                    self.detokenizer_cache[n] = {}
+                self.detokenizer_cache[n][str(input_ids)] = extended
+                return extended
 
         s = ""
         for chunk in self.chunk_out_by_special_ids(input_ids):

src/lmql/ui/playground/src/App.jsx

@@ -763,9 +763,9 @@ const ModelResultText = styled.div`
 
   div .tag-assistant {
     display: inline-block;
-    width: 65%;
     border: 1pt solid #5c5c5c;
     margin-top: 5pt;
+    margin-right: 4%;
 
     border-radius: 8pt;
     overflow: hidden;
@@ -774,7 +774,7 @@ const ModelResultText = styled.div`
 
   div .tag-user {
     display: block;
-    margin-left: 32%;
+    margin-left: 4%;
     position: relative;
     border: 1pt solid #5c5c5c;
     border-radius: 8pt;
@@ -918,8 +918,8 @@ class Truncated extends React.Component {
 
   componentDidMount() {
     this.stepper = setInterval(() => {
-      this.setState(s => Object.assign(s, { typingOffset: s.typingOffset + 4 }))
-    }, 10)
+      this.setState(s => Object.assign(s, { typingOffset: s.typingOffset + 16 }))
+    }, 5)
   }
 
   componentWillUnmount() {

src/lmql/ui/playground/src/DecoderGraph.js

@@ -567,7 +567,11 @@ export function DecoderGraph(props) {
                             target: to
                         }
                     };
-                    cy.add({group: "edges", ...e})
+                    try {
+                        cy.add({group: "edges", ...e})
+                    } catch (e) {
+                        console.error("error adding edge", e)
+                    }
                 }))
 
                 let mostLikely = layoutDecoderGraph(cy)
@@ -588,11 +592,23 @@ export function DecoderGraph(props) {
         const renderer = {
             add_result: (output) => {
                 if (output.type == "decoder-graph-state") {
-                    // persist decoder graph in local state
-                    const raw = JSON.stringify(output.data)
-                    setRawGraphData(raw)
-                    persistedState.setItem("decoder-graph", raw, onPersistedGraphChange)
-                    setCyData(output.data)
+                    setCyData(cyData => {
+                        // persist decoder graph in local state
+                        let updated = Object.assign({nodes: [], edges: []}, cyData || {})
+                        let nodes = {}
+                        updated.nodes.forEach(n => nodes[n.id] = n)
+                        output.data.nodes.forEach(n => nodes[n.id] = n)
+                        updated.nodes = Array.from(Object.values(nodes))
+                        
+                        updated.edges = Array.from([...updated.edges, ...output.data.edges])
+                        
+                        // console.log("updating", output.data.nodes.length)
+
+                        const raw = JSON.stringify(updated)
+                        setRawGraphData(raw)
+                        persistedState.queueSetItem("decoder-graph", raw, onPersistedGraphChange)
+                        return updated
+                    })
                 } else {
                     // nop in this component
                 }

src/lmql/ui/playground/src/State.js

@@ -27,6 +27,8 @@ class PersistedState {
       this.items = {}
       this.listeners = {}
       this.restore()
+
+      this.saveQueue = {}
     }
 
     persist(k) {
@@ -116,6 +118,15 @@ class PersistedState {
         })
       }
     }
+
+    queueSetItem(key, value, exclude_listener=null) {
+      if (this.saveQueue[key]) {
+        clearTimeout(this.saveQueue[key]);
+      }
+      this.saveQueue[key] = setTimeout(() => {
+        this.setItem(key, value, exclude_listener);
+      })
+    }
   }
 
 export const persistedState = new PersistedState();