Generic LLM wrapper to support chat model interface with configurable chat prompt format

[krasserm]

Update 2023-09-08

This PR now supports further models in addition to Lllama-2 chat models. See this comment for further details. The title of this PR has been updated accordingly.

Original PR description

This PR adds a generic Llama2Chat model, a wrapper for LLMs able to serve Llama-2 chat models (like LlamaCPP, HuggingFaceTextGenInference, ...). It implements BaseChatModel, converts a list of chat messages into the required Llama-2 chat prompt format and forwards the formatted prompt as str to the wrapped LLM. Usage example:

# uses a locally hosted Llama2 chat model
llm = HuggingFaceTextGenInference(
    inference_server_url="http://127.0.0.1:8080/",
    max_new_tokens=512,
    top_k=50,
    temperature=0.1,
    repetition_penalty=1.03,
)

# Wrap llm to support Llama2 chat prompt format.
# Resulting model is a chat model
model = Llama2Chat(llm=llm)

messages = [
    SystemMessage(content="You are a helpful assistant."),
    MessagesPlaceholder(variable_name="chat_history"),
    HumanMessagePromptTemplate.from_template("{text}"),
]

prompt = ChatPromptTemplate.from_messages(messages)
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
chain = LLMChain(llm=model, prompt=prompt, memory=memory)

# use chat model in a conversation
# ...

Also part of this PR are tests and a demo notebook.

• Tag maintainer: @hwchase17
• Twitter handle: @mrt1nz

[vercel]

The latest updates on your projects. Learn more about Vercel for Git ↗︎

1 Ignored Deployment

Name	Status	Preview	Comments	Updated (UTC)
langchain	⬜️ Ignored (Inspect)	Visit Preview		Nov 17, 2023 10:53pm

[krasserm]

Not included yet is support for function calls (i.e. FunctionMessage). I plan to add this in another PR.

[krasserm]

@baskaryan @rlancemartin First of all, apologies for the extra round caused by linting errors from my first commit which are now resolved by my second commit.

This was related to some issues I had running all make targets locally, for which a fix is now available in #8344. I'm still experiencing the issues mentioned in #6182 though but was able to workaround by using /usr/bin/make directly.

[rlancemartin]

@baskaryan @rlancemartin First of all, apologies for the extra round caused by linting errors from my first commit which are now resolved by my second commit.

This was related to some issues I had running all make targets locally, for which a fix is now available in #8344. I'm still experiencing the issues mentioned in #6182 though but was able to workaround by using /usr/bin/make directly.

Thanks for adding!

Re function calling, have you seen this and do you have a PR in flight to support it? Would be great!

[rlancemartin]

Nice to see that these tokens are now officially recommended, as you point out.

I was wondering whether they were going to get broadly suggested.

[krasserm]

Inspiring to see how far GPT4 gets deriving a llama2 prompt from analyzing generation.py. Looks good except from the wrongly placed first [INST].

[rlancemartin]

It would be nice to show this running w/ llama2 locally, as well.

If you don't have it locally, I can add.

[krasserm]

Good idea, I'll give this a try today or tomorrow, extend the notebook and update this PR.

[krasserm]

Thanks for reviewing my PR @rlancemartin! I haven't seen the grammar-based sampling PR, thanks for the pointer. I'll start working on it later this week and submit another PR.

[krasserm]

@rlancemartin I just added two more commits that include an example how to use the Llama2Chat wrapper with a local LlamaCpp LLM and further tests incl. failure cases. I also tried to improve the documentation in the notebook.

[krasserm]

@rlancemartin @baskaryan is there anything else you think I should add/improve to get this PR merged?

[yaoaifiling]

@krasserm The llm created by HuggingFaceTextGenInference doesn't support async operation, so it fails in async mode. Do you notice that?

[krasserm]

@yaoaifiling async support for HuggingFaceTextGenInference was recently added. Are you using an older LangChain version?

[yaoaifiling]

@yaoaifiling async support for HuggingFaceTextGenInference was recently added. Are you using an older LangChain version?

Thanks! Yes, using an old one, need to update my langchain version.

[baskaryan]

@rlancemartin @baskaryan is there anything else you think I should add/improve to get this PR merged?

apologies for the delay, will take another look shortly!

[avoroshilov]

Thank you @krasserm for doing this!
By the way, I tried it out and wanted to see how good it extends. By default it doesn't work with e.g. Vicuna, mostly because the prompt templates are different. But based on https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py#L336 I was able to make it work w/ Vicuna with very simple modifications to _to_llama2_chat_prompt - to make the template look like

{SYS} USER: {message[2n]} ASSISTANT: {message[2n+1]} </s>USER: {message[2(n+1)]} ASSISTANT: {message[2(n+1)+1]} </s>

So, this looks promising as a chat wrapper for some major local llama-based models, if templating mechanism to be modified. I also have AutoGPTQ pipeline added to LangChain locally (will probably fire a PR soon), which allows me to run quantized models easily.

How do you think this should go? This generalization is probably out of scope of this PR? If so, then we can potentially duplicate the Llama2Chat into something like LlamaLocalChat or something with generalized templating.

[krasserm]

Good thoughts @avoroshilov, thank you! I just updated this PR to make the chat wrapper generic. The new base class ChatWrapper can now be configured to support custom chat prompt formats e.g.

class Llama2Chat(ChatWrapper):
    sys_beg: str = "<s>[INST] <<SYS>>\n"
    sys_end: str = "\n<</SYS>>\n\n"
    ai_n_beg: str = " "
    ai_n_end: str = " </s>"
    usr_n_beg: str = "<s>[INST] "
    usr_n_end: str = " [/INST]"
    usr_0_beg: str = ""
    usr_0_end: str = " [/INST]"


class Llama2Instruct(ChatWrapper):
    sys_beg: str = "### System:\n"
    sys_end: str = "\n\n"
    ai_n_beg: str = "### Assistant:\n"
    ai_n_end: str = "\n\n"
    usr_n_beg: str = "### User:\n"
    usr_n_end: str = "\n\n"


class Vicuna(ChatWrapper):
    sys_beg: str = ""
    sys_end: str = " "
    ai_n_beg: str = "ASSISTANT: "
    ai_n_end: str = " </s>"
    usr_n_beg: str = "USER: "
    usr_n_end: str = " "

These classes are part of the PR (as sample implementations). Lllama2Chat is a refactoring of the original implementation. Llama2Instruct is a wrapper for llama-2-*-instruct-v2 models (and also several other Llama-2 derivatives like StableBeluga-2) that require a prompt format like

### System:
{System}

### User:
{User}

### Assistant:
{Assistant}

Vicuna is for the prompt format you posted i.e.

{SYS} USER: {message[2n]} ASSISTANT: {message[2n+1]} </s>USER: {message[2(n+1)]} ASSISTANT: {message[2(n+1)+1]} </s>

I hope this fits your use cases!

[krasserm]

@baskaryan @rlancemartin it would be great if you could take another look at this PR. Based on @avoroshilov's thoughts I made this a generic LLM wrapper that now supports a lot more models than just Llama-2 chat models. Models not already covered can now be easily supported by extending ChatWrapper base class. See this comment details.

[cstub]

@krasserm This looks fantastic! I've been looking for an integration of Llama-2 based chat models in LangChain. Would love to see this merged soon! 👍

[shaw91]

Hi team, any update on this PR? I don't see llama-2-based chat model support in LangChain yet. So being able to merge this would be really awesome.

[leo-gan]

@krasserm Hi , could you, please, resolve the merging issues and address the last comments (if needed)? After that, ping me and I push this PR for the review. Thanks!

[krasserm]

@leo-gan thanks for your support! I just rebased on master and made some changes to the notebook (support for latest llama-cpp-python version, ...). I also squashed all commits from this PR.

[leo-gan]

@baskaryan Please review this PR. TNX

[maziyarpanahi]

It would be very helpful to the community using foundational LLMs if we can have this PR merged.

[efriis]

Hey folks! Appreciate your patience. Given this is effectively a model-specific prompt template, I would prefer to have this as a reference prompt at smith.langchain.com/hub , or as a LangChain template for working with open source models instead of a "model" as it's implemented here.

If nobody else picks it up, I can spend some time putting together a little template with alternatives showing how to format this, so people can get started with these llms quicker with that, without merging these prompts into the models codebase.

[maziyarpanahi]

Hey folks! Appreciate your patience. Given this is effectively a model-specific prompt template, I would prefer to have this as a reference prompt at smith.langchain.com/hub , or as a LangChain template for working with open source models instead of a "model" as it's implemented here.

If nobody else picks it up, I can spend some time putting together a little template with alternatives showing how to format this, so people can get started with these llms quicker with that, without merging these prompts into the models codebase.

Hi @efriis

I guess this PR could have changed in a way to not call on each LLM model, but any LLM model that requires a specific start/end for system, user, ai, etc messages. However, if we are not going to merge this, I am interested in LangChain template for working with open source models solution which doesn't require downloading something extra from a Hub or installing another dependency just to get something like Llama-2 work.

[krasserm]

I decided for this design because I see start/end tokens in chat prompts as an implementation detail of a model. This is something a user shouldn't need to care about when designing an application-specific chat prompt template like the one used in the initial example e.g.

messages = [
    SystemMessage(content="You are a helpful assistant."),
    MessagesPlaceholder(variable_name="chat_history"),
    HumanMessagePromptTemplate.from_template("{text}"),
]

# independent of model-specific chat message start/end tokens
prompt_template = ChatPromptTemplate.from_messages(messages)

Start/end tokens for messages of a chat model are an implementation detail in the same way as a tokenizer is a model-specific implementation detail. I think that's the reason why Hugging Face decided to link chat templates to tokenizers. Their chat templates too render model-independent chat messages into prompts with model-specific start/end tokens (like the wrappers of this PR) and are reusable across tokenizers.

You can also see the wrapper(s) from this PR as a layer in a chat protocol, translating messages of a conversation into a lower-level, model-specific prompt format, freeing the user from dealing with these details. And these wrappers (or protocol layers) are reusable too. For example, you can reuse the Llama2Chat wrapper with all LLMs that require the Llama-2 chat prompt (and many of them do, not only vanilla Llama-2 chat models).

Even if you do not want to have concrete subclasses, like Llama2Chat, part of this PR, I think the ChatWrapper base class is still a useful addition to the models codebase as it provides the generic functionality for implementing a uniform chat model interface for many LangChain LLMs. For example, many models require the ChatML prompt format which can be implemented with ChatWrapper easily like:

class ChatML(ChatWrapper):
    sys_beg: str = "<|im_start|>system\n"
    sys_end: str = "\n<|im_end|>"
    ai_n_beg: str = "<|im_start|>assistant\n"
    ai_n_end: str = "\n<|im_end|>"
    usr_n_beg: str = "<|im_start|>user\n"
    usr_n_end: str = "\n<|im_end|>"

Their wide applicability is one of the main reasons why I thought concrete ChatWrapper implementations should be part of this PR.

[efriis]

@krasserm thoughts on merging this into experimental initially? Code can stay almost the same - all the dependencies would still come from langchain, but the new stuff would come from langchain_experimental while we ask people to try it out and let us know their thoughts, and how it plays with the different models.

[krasserm]

@efriis sounds good to me. I'll update the PR later this week.

[efriis]

Sounds good. Thanks!

Another (better imo) implementation is as some kind of function that formats ChatPrompts into the format expected by these LLMs, and an output parser that parses the output into a message (or messages).

I would imagine a runnable like

chain = ChatPromptTemplate.from_message([...]) | new_prompt_formatter | HuggingFaceTextGenInterface(...) | new_llama_output_parser 

[maziyarpanahi]

Sounds good. Thanks!

Another (better imo) implementation is as some kind of function that formats ChatPrompts into the format expected by these LLMs, and an output parser that parses the output into a message (or messages).

I would imagine a runnable like

chain = ChatPromptTemplate.from_message([...]) | new_prompt_formatter | HuggingFaceTextGenInterface(...) | new_llama_output_parser 

It would be great if that ChatPromptTemplate can have a format that can be applied to every chain, agent, etc. that uses any ChatPrompt messages. This does override the ChatPromptValue for everything else #9917

[krasserm]

Sounds good. Thanks!

Another (better imo) implementation is as some kind of function that formats ChatPrompts into the format expected by these LLMs, and an output parser that parses the output into a message (or messages).

I would imagine a runnable like

chain = ChatPromptTemplate.from_message([...]) | new_prompt_formatter | HuggingFaceTextGenInterface(...) | new_llama_output_parser 

So you're basically creating a chain that allows using an LLM as a chat model without using an implementation of BaseChatModel. What is the purpose of the BaseChatModel asbtraction then?

From my perspective, BaseChatModel implementations should also be responsible for hiding details like serializing chat messages into low-level prompt strings. Users of BaseChatModel implementations don't (and shouldn't) need to care about these details.

For example, when using ChatOpenAI, details like separating chat messages within a lower-level prompt string are also hidden from the user (it is even hidden behind the OpenAI API but that's an implementation detail too).

So why shouldn't LangChain try to follow these design principles of chat model interfaces consistently i.e. encapsulating details like chat messages separators, tokenization, ... etc. behind chat model interfaces instead of sometimes exposing it to users (like in your proposal) and sometimes not?

Having a common chat model interface with clear semantics (which includes dealing with low-level chat message separators) would simplifiy application development with LangChain a lot.

[krasserm]

I just moved everything from langchain to experimental and updated the demo notebook accordingly. I also added a pytest-asyncio dependency to pyproject.toml needed for executing async tests. An updated poetry.lock is included. Commits are rebased on master.

[efriis]

Appreciate it!

Re: the previous point about composing the LLM rather than wrapping it in the BaseChatModel - I think this is a discussion of simplicity vs ease of use.

ChatOpenAI makes sense as a purpose-built chat interface because it consistently returns the correct format. With many of these self-hosted LLMs, getting that chat model kind of consistency is difficult, so having simpler abstractions that have to be composed together makes more sense to me.

That being said, these kinds of experiments are what experimental is for! Thanks for moving it

[efriis]

Landed! Thanks @krasserm and appreciate your patience on this one

[krasserm]

Thanks @efriis for discussing and supporting this PR