[ChatStream] Llama2 対応の ChatPrompt実装

こんにちは！ (株)Qualiteg プロダクト開発部 です！

本稿では、 ChatStream にLlama2 対応の ChatPrompt を同梱いたしましたのでご紹介いたします！

現在の ChatPrompt は以下とおりとなっております。旧バージョンのChatStreamをご利用の場合も、以下コードにて Llama2対応可能です。（もちろん最新バージョンのChatStreamには同梱済です）

from chatstream import AbstractChatPrompt
from chatstream.chat_prompt.role_type import RoleType

SYSTEM_PROMPT = """\
You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.\
"""


class ChatPromptMetaLlamaLlama2Chat(AbstractChatPrompt):
    """
    meta-llama/Llama-2-7b-chat

    Prompt Guide from
    https://huggingface.co/blog/llama2
    """

    def __init__(self):
        super().__init__()  # Call the initialization of the base class
        self.set_system(f"<s>[INST] <<SYS>>\n{SYSTEM_PROMPT}\n<</SYS>>\n\n")
        self.set_requester("")
        self.set_responder("")

    def get_stop_strs(self):
        if not self.chat_mode:
            return None
        return []

    def get_custom_skip_echo_len(self, skip_echo_len):
        # modify skip_echo_len when using llama2
        num_turn = self.get_turn()
        if num_turn >= 2:
            modified_skip_echo_len = skip_echo_len + 1 * self.get_turn()
            return modified_skip_echo_len
        return skip_echo_len

    def get_replacement_when_input(self):
        return None

    def get_replacement_when_output(self):  # replace when response_text gotten
        return None

    def create_prompt(self, opts={}):
        if self.chat_mode == False:
            return self.get_requester_last_msg()

        # Chat Mode == True の場合のプロンプトを構築する
        ret = self.system

        for chat_content in self.get_contents(opts):

            chat_content_role_type = chat_content.get_role_type()
            chat_content_message = chat_content.get_message()

            if chat_content_message:
                merged_message = ""
                if chat_content_role_type == RoleType.REQUESTER:
                    merged_message = f"{chat_content_message} [/INST] "
                elif chat_content_role_type == RoleType.RESPONDER:
                    merged_message = f"{chat_content_message} </s><s>[INST] "
                ret += merged_message
            else:
                pass

        return ret

    def build_initial_prompt(self, chat_prompt):
        # 初期プロンプトは実装しない
        pass

本実装については、１点だけトリッキーな実装を含んでいます。
というのも、ChatStreamはストリーミングチャットが基本ですので新たに生成された１トークン分を更新していくことでリアルタイムで流れるようなチャットを実現しています。そのため、出力が期待と１文字でもズレてしまうと、ストリーミング時に１文字かけたり、全体が１文字ずつズレてしまったりする現象が発生しやすいです。

Llama2の出力の場合、ターンごとに1文字ずつ出力に空白文字が発生する現象が確認されており、 なにがしかのSpecial Token や空トークン、またコード変換の影響を疑っていましたが根本原因の特定に至らず、現象をキャンセルするための実装を入れています。

現象としましては、

以下のプロンプトを入力プロンプトとしたとき、

<s>[INST] <<SYS>>
You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
<</SYS>>

Do you know the movie titanic [/INST] Hello! Yes, I'm familiar with the movie Titanic. It's a classic film directed by James Cameron, released in 1997, and starring Leonardo DiCaprio and Kate Winslet. The movie is based on the true story of the RMS Titanic, a British passenger liner that sank in the North Atlantic Ocean in 1912 after colliding with an iceberg. The film follows the story of Jack Dawson (played by DiCaprio) and Rose DeWitt Bukater (played by Winslet), who come from </s><s>[INST] Who is the director [/INST] The director of the movie "Titanic" is James Cameron. </s><s>[INST] Who is starred [/INST] ]  The movie "Titanic" features a star-studded cast, including:
* Leonardo DiCaprio as Jack Dawson
* Kate Winslet as Rose DeWitt Bukater
* Billy Zane as Cal Hockley
* Kathy Bates as Molly Brown
* Frances Fisher as Ruth DeWitt Bukater
* Bernard Hill as Captain Edward John Smith
* Jonathan Hyde as J. Bruce Ismay
* Eric Braeden as John Jacob Astor IV
* Gloria Stuart as Old Rose

These actors brought the characters from the movie to life and </s><s>[INST] 

出力は以下のようになります

[INST] <<SYS>>
You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
<</SYS>>

Do you know the movie titanic [/INST] Hello! Yes, I'm familiar with the movie Titanic. It's a classic film directed by James Cameron, released in 1997, and starring Leonardo DiCaprio and Kate Winslet. The movie is based on the true story of the RMS Titanic, a British passenger liner that sank in the North Atlantic Ocean in 1912 after colliding with an iceberg. The film follows the story of Jack Dawson (played by DiCaprio) and Rose DeWitt Bukater (played by Winslet), who come from  [INST] Who is the director [/INST] The director of the movie "Titanic" is James Cameron.  [INST] Who is starred [/INST]

Input:
Do you know the movie titanic [/INST] Hello! Yes, I'm familiar with the movie Titanic. It's a classic film directed by James Cameron, released in 1997, and starring Leonardo DiCaprio and Kate Winslet. The movie is based on the true story of the RMS Titanic, a British passenger liner that sank in the North Atlantic Ocean in 1912 after colliding with an iceberg. The film follows the story of Jack Dawson (played by DiCaprio) and Rose DeWitt Bukater (played by Winslet), who come from [INST] Who is the director [/INST] The director of the movie "Titanic" is James Cameron. [INST] Who is starred [/INST] 

つまり入力プロンプトが [INST] つまり 半角スペース１つ ＋ [INST]　でも出力で、 [INST]＝ 半角スペース２つ ＋ [INST]　となって戻ってくるため、新規生成文章の切り出しで１レスポンスにたいして１文字ズレの発生が観測されます。

それならばと、入力プロンプトを 半角スペース2つ ＋ [INST] で入力したが、その場合は出力生成文は 半角スペース3つ ＋[INST]
となり、ズレが解消されないようです。

このとき、通常のテキスト処理で trim していく方法もありますが、当ChatPromptではなるべくもとの入力と出力を変形しないことをポリシーとして、

    def get_custom_skip_echo_len(self, skip_echo_len):
        # modify skip_echo_len when using llama2
        num_turn = self.get_turn()
        if num_turn >= 2:
            modified_skip_echo_len = skip_echo_len + 1 * self.get_turn()
            return modified_skip_echo_len
        return skip_echo_len

のように、 skip_echo_len をターンにあわせて調整することでピッタリの切り出しに対応しました。約１０００通りの入出力において、問題なく切り出せているため、本手法を workaround として採用しています。ちょっとした見落としの可能性も高そうなので根治のために研究部側で引き続き原因調査しています。

Llama2 をベースモデルとして採用しているほかモデルも同様の現象が発生していますが、本手法にて解決できました。