Terminal output color and streaming!

[dimondjik]

Here's a code snippet you guys may find helpful.
(guidance >= 0.1 since it's the version where streaming was introduced, python 3.10.6)
It outputs colored stream into terminal (I personally use it in PyCharm), green is the text generated by llm:

from guidance import guidance, gen, models
from llama_cpp import Llama
from html.parser import HTMLParser


# Custom class based on in-built python html parser
class ModelStateHTMLParser(HTMLParser):
    def __init__(self):
        super().__init__()
        self.colored_text = ''

    def feed(self, data):
        self.colored_text = ''
        # Remove html insertion tags (I suppose it is for jupyter to recognize it as html markdown)
        data = data.replace('<||_html:', '').replace('_||>', '')
        super().feed(data)
        return self.colored_text
   
    # Start ASCII coloring text green (32) when span tag opens and reset color (0) when tag closes
    def handle_starttag(self, tag, attrs):
        if tag == 'span':
            self.colored_text += '\033[32m'

    def handle_endtag(self, tag):
        if tag == 'span':
            self.colored_text += '\033[0m'

    def handle_data(self, data):
        self.colored_text += data


@guidance(stateless=False)
def test(lm):
    return lm + 'This is an article' + gen(stop='.', suffix='. ') + About that great library called "guidance"' + gen(stop='\n')


llm_llamacpp = Llama(model_path="models/mistral-7b-instruct-v0.2.Q5_K_M.gguf",
                     n_gpu_layers=-1, n_threads=32, n_ctx=4096, verbose=False)

llm = models.LlamaCpp(llm_llamacpp)

state_parser = ModelStateHTMLParser()
prev_part_len = 0

# Thanks for this convenient stream!
for part in llm.stream() + test():
    # Assess private state field of the model and feed it to parser, cutting out symbols we already printed
    print(state_parser.feed(part._state[prev_part_len:]), end='')
    prev_part_len = len(part._state)
print()

[Harsha-Nori]

Thanks for sharing, this is super cool! I think if we can accurately detect when users are on a terminal (vs ipython), this would be a great utility to include in the package. Let me think on if we can detect this reasonably well...thanks for sharing again!

[hudson-ai]

+1 super, super cool. I would love full-featured terminal support :)

[hudson-ai]

@dimondjik I hope you don't mind that I incorporated this code into a PR -- please let me know if you have any feedback or suggestions :)