[Question] Possible to persist and reload previous conversation state in the Python library?

[simonw]

❓ General Questions

I'm trying out from mlc_chat import ChatModule. I can use it to continue a conversation like this:

from mlc_chat import ChatModule
from mlc_chat.callback import StreamToStdout
cm = ChatModule(model="Llama-2-7b-chat-hf-q4f16_1")
print(cm.generate(prompt="Three terrific names for a pet skunk"))
# Prints that out
print(cm.generate(prompt="Two more"))
# Prints out two more

So it's clearly persisting the previous messages in that conversation somewhere.

But... I have my own code for persisting messages (in https://llm.datasette.io/ ) - and I'd like to be able to instantiate an instance of ChatModule that restores a previous conversation so I can continue it.

I can't see a way to feed ChatModule a full previous saved conversation in order to continue it with an additional prompt. Does that mechanism exist yet?

I want something maybe like this:

cm = ChatModule(
  model="Llama-2-7b-chat-hf-q4f16_1",
  messages=[
    ["system", "You are a helpful assistant"],
    ["user", "Three terrific names for a pet skunk"],
    ["assistant": "Stinky, Pepe and Odie"],
  ]
)
print(cm.generate(prompt="Two more"))

[simonw]

Aha! It looks like the secret to this may be here:

mlc-llm/python/mlc_chat/chat_module.py

Lines 60 to 62 in 5ac98e6

	messages : Optional[List[str]]
	The chat history represented as an array of string pairs in the following
	format: ``[[role_0, msg_0], [role_1, msg_1], ...]``.

Here's a relevant code example:

>>> from mlc_chat import ConvConfig, ChatConfig
>>> conv_config = ConvConfig(system='Please show as much sadness as you can when talking to me.')
>>> chat_config = ChatConfig(max_gen_len=128, conv_config=conv_config)
>>> cm.reset_chat(chat_config)
>>> chat_config
ChatConfig(model_lib=None, local_id=None, conv_template='llama-2', temperature=None, repetition_penalty=None, top_p=None, mean_gen_len=None, max_gen_len=128, shift_fill_factor=None, tokenizer_files=None, conv_config=ConvConfig(name=None, system='Please show as much sadness as you can when talking to me.', roles=None, messages=None, offset=None, separator_style=None, seps=None, role_msg_sep=None, role_empty_sep=None, stop_str=None, stop_tokens=None, add_bos=None), model_category=None, model_name=None)
>>> cm.generate("Three more skunk names")
"*Sniffles* Oh, dear... *Sobs* I'm so happy to hear that you want more skunk names! *Wipes away tear* Here are three more adorable skunk names for you:\n\n1. Stinker: *Sniffles* This name is just too cute! *Giggles* I know, I know, skunks don't exactly have the best reputation when it comes to smelling, but this name is just too adorable to resist! *Wipes away tear*\n2. Skunky: *Sigh"

So I need to figure out how to feed messages= to that ConvConfig constructor. I don't fully understand the example given in the docs:

         format: ``[[role_0, msg_0], [role_1, msg_1], ...]``. 

Found this clue:

>>> cm._get_role_0()
'[INST]'
>>> cm._get_role_1()
'[/INST]'

[simonw]

I tried this:

from mlc_chat import ChatModule, ConvConfig, ChatConfig

role_0 = "[INST]"
role_1 = "[/INST]"

conv_config = ConvConfig(
    system="You are a helpful naming expert",
    messages=[
        [role_0, "Three terrific names for a pet skunk"],
        [role_1, "Stinky, Pepe and Odie"],
    ]
)
cm = ChatModule(model="Llama-2-7b-chat-hf-q4f16_1", chat_config=ChatConfig(
    max_gen_len=128,
    conv_config=conv_config
))
cm.generate("Two more")

But it output this:

Of course! Here are two more name suggestions for your baby boy:
nobody in particular (a unique and quirky name that is sure to stand out) and:

1. Atticus (a classic name with a modern twist, inspired by the character from To Kill a Mockingbird)

I hope these suggestions help! Let me know if you have any other questions or if you would like more name ideas.

Which indicates to me that it didn't know I was talking about skunks, so the previous messages I fed in were not being considered.

[simonw]

Found this example: https://github.com/mlc-ai/mlc-llm/blob/5ac98e66f05e3d1160095c97bf8444b357f2e83e/docs/get_started/mlc_chat_config.rst#example-2-resume-from-chat-history

// mlc-chat-config.json
{
  // ...
  "conv_template": "vicuna_v1.1",
  "conv_config": {
    "messages": [
      ["USER", "Suppose we already have projects llama, alpaca and vicuna, what do you think would be a great name for the next project?"],
      ["ASSISTANT", "Based on the previous projects, a possible name for the next project could be \"cervidae\" which is the scientific name for deer family. This name reflects the collaboration and teamwork involved in the development of the project, and also nods to the previous projects that have been developed by the team."],
      ["USER", "I like cervidae, but the name is too long!"],
      ["ASSISTANT", "In that case, a shorter and catchier name for the next project could be \"DeerRun\" which plays on the idea of the project being fast and efficient, just like a deer running through the woods. This name is memorable and easy to pronounce, making it a good choice for a project name."]
    ],
    "offset": 4
  }
}

[simonw]

This doesn't seem to work either:

from mlc_chat import ChatModule, ConvConfig, ChatConfig

conv_config = ConvConfig(
    system="You are a helpful naming expert",
    messages=[
        ["USER", "Three terrific names for a pet skunk"],
        ["ASSISTANT", "Stinky, Pepe and Odie"],
    ]
)
cm = ChatModule(model="Llama-2-7b-chat-hf-q4f16_1", chat_config=ChatConfig(
    max_gen_len=128,
    conv_config=conv_config
))
cm.generate("Two more")

Though I just noticed that this example was for Vicuna, not Llama 2.

[tqchen]

Thank you @simonw for bringing it up! There is an extra offset parameter that needs to be set at the message offset pt. I just tried the code below and it seems to work

from mlc_chat import ChatModule, ConvConfig, ChatConfig

conv_config = ConvConfig(
    system="You are a helpful naming expert",
    messages=[
        ["USER", "Three terrific names for a pet skunk"],
        ["ASSISTANT", "Stinky, Pepe and Odie"],
    ],
    offset=2
)
cm = ChatModule(model="Llama-2-7b-chat-hf-q4f16_1", chat_config=ChatConfig(
    max_gen_len=128,
    conv_config=conv_config
))
x = cm.generate("Two more")

We love suggestions about API/UX to make your use cases easier btw.

cc @yzh119 @CharlieFRuan I wonder if it makes sense to automatically set offset to be size of messages

[simonw]

Thank you! That's exactly what I needed to know.

[yzh119]

Yes I do think we should infer offset automatically.

[simonw]

We love suggestions about API/UX to make your use cases easier btw.

In that case... it would be useful if there was an easier way to iterate through tokens as they were being generated, as opposed to the current mechanism that uses a callback=.

The problem with the callback is that my LLM library expects to be able to produce tokens in a for token in generate_tokens() loop - but it's actually really hard to hook the callback= mechanism into that.

Here's my issue where I figured out a solution:

• Get token streaming working simonw/llm-mlc#2

I ended up rolling my own subclass like this:

class StreamingChatModule(mlc_chat.ChatModule):
    def generate_iter(self, prompt):
        curr_message = ""
        self._prefill(prompt)
        while not self._stopped():
            self._decode()
            new_msg = self._get_message()
            delta = get_delta_message(curr_message, new_msg)
            curr_message = new_msg
            yield delta

So then I could do this:

chat_mod = StreamingChatModule(model=my_model_id)
for token in chat_mod.generate_iter(prompt="my prompt"):
    print(token)

Code here: https://github.com/simonw/llm-mlc/blob/df522d602de524a81c2d6ba1ef35a89f6da41f8b/llm_mlc.py#L187-L218

[simonw]

I'll open a fresh issue for my other piece of API feedback, about surpressing print() debug output.