Convert the syntax to use with blocks, and fix some context block ordering bugs

Author: slundberg

guidance/library/__init__.py

@@ -13,9 +13,10 @@
 
 # context blocks
 from ._block import block
-from ._role import role, system, assistant, user, function, instruction
+from ._role import role, system, assistant, user, function, instruction, indent_roles
 from ._format import monospace
 from ._silent import silent
+from ._set_var import set_var
 # from ..models._model import context_free
 
 # stateless library functions

guidance/library/_role.py

@@ -1,5 +1,6 @@
 import guidance
 from ._block import block
+from ._set_var import set_var
 
 nodisp_start = "<||_#NODISP_||>"
 nodisp_end = "<||_/NODISP_||>"
@@ -8,50 +9,54 @@
 
 
 @guidance
-def role_opener(lm, role_name, debug=False, **kwargs):
+def role_opener(lm, role_name, **kwargs):
+    indent = lm.get("__role_indent", True)
     if not hasattr(lm, "get_role_start"):
         raise Exception(
             f"You need to use a chat model in order the use role blocks like `with {role_name}():`! Perhaps you meant to use the {type(lm).__name__}Chat class?"
         )
 
     # Block start container (centers elements)
-    lm += f"<||_html:<div style='display: flex; border-bottom: 1px solid rgba(127, 127, 127, 0.2);  justify-content: center; align-items: center;'><div style='flex: 0 0 80px; opacity: 0.5;'>{role_name.lower()}</div><div style='flex-grow: 1; padding: 5px; padding-top: 10px; padding-bottom: 10px; margin-top: 0px; white-space: pre-wrap; margin-bottom: 0px;'>_||>"
+    if indent:
+        lm += f"<||_html:<div style='display: flex; border-bottom: 1px solid rgba(127, 127, 127, 0.2);  justify-content: center; align-items: center;'><div style='flex: 0 0 80px; opacity: 0.5;'>{role_name.lower()}</div><div style='flex-grow: 1; padding: 5px; padding-top: 10px; padding-bottom: 10px; margin-top: 0px; white-space: pre-wrap; margin-bottom: 0px;'>_||>"
 
     # Start of either debug or HTML no disp block
-    if debug:
-        lm += span_start
-    else:
+    if indent:
         lm += nodisp_start
+    else:
+        lm += span_start
 
     lm += lm.get_role_start(role_name, **kwargs)
 
     # End of either debug or HTML no disp block
-    if debug:
-        lm += span_end
-    else:
+    if indent:
         lm += nodisp_end
+    else:
+        lm += span_end
 
     return lm
 
 
 @guidance
-def role_closer(lm, role_name, debug=False, **kwargs):
+def role_closer(lm, role_name, **kwargs):
+    indent = lm.get("__role_indent", True)
     # Start of either debug or HTML no disp block
-    if debug:
-        lm += span_start
-    else:
+    if indent:
         lm += nodisp_start
+    else:
+        lm += span_start
 
     lm += lm.get_role_end(role_name)
 
     # End of either debug or HTML no disp block
-    if debug:
-        lm += span_end
-    else:
+    if indent:
         lm += nodisp_end
+    else:
+        lm += span_end
 
     # End of top container
-    lm += "<||_html:</div></div>_||>"
+    if indent:
+        lm += "<||_html:</div></div>_||>"
 
     return lm
 
@@ -85,3 +90,6 @@ def function(text=None, **kwargs):
 
 def instruction(text=None, **kwargs):
     return role("instruction", text, **kwargs)
+
+def indent_roles(indent=True):
+    return set_var("__role_indent", indent)

guidance/library/_set_var.py

@@ -0,0 +1,21 @@
+import guidance
+from ._block import block
+
+@guidance
+def set_opener(lm, name, value):
+    if name in lm:
+        lm = lm.set("__save" + name, lm[name])
+    return lm.set(name, value)
+
+@guidance
+def set_closer(lm, name):
+    if "__save" + name in lm:
+        return lm.set(name, lm["__save" + name]).remove("__save" + name)
+    else:
+        return lm.remove(name)
+
+def set_var(name, value=True):
+    return block(
+        opener=set_opener(name, value),
+        closer=set_closer(name),
+    )

guidance/models/_model.py

@@ -252,31 +252,44 @@ def __add__(self, value):
         # inside this context we are free to drop display calls that come too close together
         with throttle_refresh():
 
-            # close any newly closed contexts
+            # find what new blocks need to be applied
+            new_blocks = []
+            for context in Model.open_blocks:
+                if context not in lm.opened_blocks:
+                    new_blocks.append(context)
+
+                    # mark this so we don't re-add when computing the opener or closer (even though we don't know the close text yet)
+                    lm.opened_blocks[context] = (0, "")
+
+            # find what old blocks need to be removed
+            old_blocks = []
             for context in list(reversed(lm.opened_blocks)):
                 if context not in Model.open_blocks and context in lm.opened_blocks:
-                    pos, close_text = lm.opened_blocks[context] # save so we can delete it before adding it
-                    if context.name is not None:
-                        lm._variables[context.name] = format_pattern.sub("", lm._state[pos:])
+                    old_blocks.append((lm.opened_blocks[context], context))
+
+                    # delete this so we don't re-close when computing the opener or closer
                     del lm.opened_blocks[context]
-                    lm._inplace_append(close_text)
+
+            # close any newly closed contexts
+            for (pos, close_text), context in old_blocks:
+                if context.name is not None:
+                    lm._variables[context.name] = format_pattern.sub("", lm._state[pos:])
+                lm += context.closer
 
             # apply any newly opened contexts (new from this object's perspective)
-            for context in Model.open_blocks:
-                if context not in lm.opened_blocks:
-                    lm.opened_blocks[context] = (0, "") # mark this so we don't readd when computing the opener (even though we don't know the close text yet)
-                    lm += context.opener
-                    with grammar_only():
-                        tmp = lm + context.closer
-                    close_text = tmp._state[len(lm._state):] # get the new state added by calling the closer
-                    lm.opened_blocks[context] = (len(lm._state), close_text)
-                    
-                    # clear out names that we override
-                    if context.name is not None:
-                        if context.name in lm._variables:
-                            del lm._variables[context.name]
-                            if context.name in lm._variables_log_probs:
-                                del lm._variables_log_probs[context.name]
+            for context in new_blocks:
+                lm += context.opener
+                with grammar_only():
+                    tmp = lm + context.closer
+                close_text = tmp._state[len(lm._state):] # get the new state added by calling the closer
+                lm.opened_blocks[context] = (len(lm._state), close_text)
+                
+                # clear out names that we override
+                if context.name is not None:
+                    if context.name in lm._variables:
+                        del lm._variables[context.name]
+                        if context.name in lm._variables_log_probs:
+                            del lm._variables_log_probs[context.name]
 
             # wrap raw string values
             if isinstance(value, str):
@@ -957,9 +970,7 @@ def __call__(self, grammar, max_tokens=1000000, n=1, top_p=1, temperature=0.0, e
                 #     self._cache_state["new_token_ids"].append(sampled_token_ind)
 
                 # capture the named groups from the parse tree
-                new_captured_data, new_captured_log_prob_data = parser.get_captures()
-                captured_data.update(new_captured_data)
-                captured_log_prob_data.update(new_captured_log_prob_data)
+                parser.get_captures(captured_data, captured_log_prob_data)
 
                 # we have no valid log prob data if we didn't compute it
                 yield new_bytes[hidden_count:], is_generated, new_bytes_prob, captured_data, captured_log_prob_data, token_count - last_token_count

guidance/models/transformers/__init__.py

@@ -1 +1,2 @@
-from ._llama import Llama, LlamaChat
+from ._llama import Llama, LlamaChat
+from ._transformers import Transformers, TransformersChat