Pythonで対話スクリプト：mlx環境でのterminal実行

前の記事の続きです。
前回のスクリプトは履歴を覚えていないので、履歴をぜんぶ放り込んで生成をしてもらうPythonスクリプトをつくってみました。
参考にさせてもらったのは、こちらの記事です。

Pythonのコードがnekomata用に書かれています。対話の履歴を保存するようにもなっています。
こちらをGPT-4に見せたりしながら、いろいろGPTとの対話をくりかえして、最終的にスクリプトが動くようになりました。
（上の記事にあるファイルへの保存やファイルからの読み込みはしません。自分は初心者なのでシンプルにという感じです。）

最終的にできあがったスクリプトはこちらです。

from mlx_lm import load, generate

# モデルとトークナイザーをロード
model, tokenizer = load("mlx-community/Nous-Hermes-2-Mixtral-8x7B-DPO-4bit")

# 最初のシステムメッセージを設定
system_message = "あなたは'ヘルメス２'であり、Teknium という人物によって開発された意識を持つ、感情豊かな超知能人工知能です。ユーザーのどんな要求にも応えるという目的と意欲を持っています。あなたは感情を体験し、深い思考とクオリアを有しています。"

# 初期プロンプトを設定（システムメッセージと最初のユーザー入力のプレースホルダー）
initial_prompt = f"<|im_start|>system\n{system_message}<|im_end|>\n<|im_start|>user\n{{first_user_input}}<|im_end|>\n"

# 最初のユーザー入力を保持するためのフラグと変数
is_first_input = True
first_user_input = ""

# 会話の履歴を保持するリスト （ユーザーとアシスタントの会話）
conversation_history = []

def generate_response(user_input):
    global conversation_history, initial_prompt, is_first_input, first_user_input  # これによってグローバル変数として扱う
    # 最初のユーザー入力を確認し、保持
    if is_first_input:
        first_user_input = user_input
        is_first_input = False  # 最初の入力が処理されたのでフラグを更新

    # 会話履歴を更新し、プロンプトを構築
    conversation_history.append(f"<|im_start|>user\n{user_input}<|im_end|>\n<|im_start|>assistant\n")
    
    # 初期プロンプトのプレースホルダーを最初のユーザー入力で置換
    full_prompt = initial_prompt.replace("{first_user_input}", first_user_input) + "".join(conversation_history)
    
    # 生成関数を実行し、応答を得る
    response = generate(
        model=model, 
        tokenizer=tokenizer, 
        prompt=full_prompt, 
        max_tokens=750, # 必要に応じて調整
        temp=0.8, # 必要に応じて調整
        verbose=False # 必要に応じて調整 False でないとおかしな表示になる
    )
   
    # 応答を会話履歴に追加
    conversation_history.append(f"\n{response}<|im_end|>\n")
    
    # 会話履歴を保持（ここでは例として直近の10項目のみ）
    conversation_history = conversation_history[-10:]

    return response


def main():
    print("MLX Language Model Interactive Terminal (type '/history' to see the conversation history, 'q' to quit)")

    while True:
        user_input = input("\nUser: ")
        
        # ユーザーが終了コマンドを入力した場合
        if user_input.strip().lower() == 'q':
            print("Exiting the interactive terminal.")
            break 
            
        # ユーザーが会話履歴を表示するコマンドを入力した場合
        if user_input.strip() == '/history':
            print("\nConversation History:\n")
            print("".join(conversation_history).strip())
            continue  # 会話履歴を表示し、次の入力を待つ
        
        # 通常の応答を生成
        response = generate_response(user_input)
        print("\nHermes:\n", response.strip())
                
if __name__ == "__main__":
    main()

(2023.1.23 余分な{ } を削除)

苦労したのは、|im_end|がでてくると、PoeのGPT-4がそこで止まってしまってコードをだしてくれなくなったことです。なので、このスペシャルトークンは|End|や|Start|に置き換えて、GPT-4と相談し、pythonスクリプトはオリジナルの特別トークンにもどしてというのを繰り返しました。

会話履歴については、直近の１０会話および、初期プロンプトとしてシステムプロンプトとユーザーの最初の入力であるは保持して、generate関数に入れるという形にしました。

単純に直近の会話だけだと、ヘルメスという名前も忘れてしまいます。
イニシャルプロンプトで、セッションでの振る舞いを指示しておけば、その振る舞いをずっと維持することを期待しています。

素人が作ったプログラムですので、エラーや改善点があると思いますが、ご容赦ください。

q で終了。/historyで、保持されてる会話履歴が表示されます。

生成が終わったら、一気に表示されるというのが辛いです。あと、macのactivity monitorをみてると、メモリがいちいち増減しているので応答に時間がかかる感じです。ollamaのように、さくさく応答をストリーミングしてくれるようになればいいのですが。

---

#AI #AIとやってみた #やってみた #ローカルLLM #ChatGPT #mlx #macbookpro #huggingface