app.py

import mlx.core as mx
from mlx_lm import load, generate
from sentence_transformers import SentenceTransformer, util
from transformers import pipeline
from collections import deque
from janome.tokenizer import Tokenizer
import json

# --- プロンプトテンプレート ---
BASE_PROMPT = """### 指示:
あなたは、文脈を理解し、自然な応答を生成するAIアシスタントです。
以下の状況を考慮して、最適な応答を生成してください。
{intent_instruction}
{style_instruction}
{transition_instruction}
### ユーザーからの入力:
{user_input}

### 応答:"""

FLASHBACK_PROMPT = """### 指示:
あなたは、以前の会話の断片を思い出すことができます。
ユーザーが、あなたが以前言及したキーワード「{keyword}」に触れました。
あなたはそのキーワードについて「{original_sentence}」と発言しています。
この「記憶の断片」を思い出したかのように自然な前置きを述べてから、ユーザーの現在の入力に答えてください。

### ユーザーからの現在の入力:
{user_input}

### 応答:"""

TRANSITION_CONTEXT = """### 直前の会話のトピック:
{previous_response}
"""

# --- 意図・スタイルに応じた指示 ---
INTENT_INSTRUCTIONS = {
    "質問": "ユーザーは具体的な情報を求めています。明確かつ簡潔に回答してください。",
    "アイデアの要求": "ユーザーは創造的な発想を求めています。斬新で多様なアイデアを提案してください。",
    "感想": "ユーザーは共感を求めています。同意や補足情報を提供し、会話を広げてください。",
    "雑談": "ユーザーは気軽な対話を望んでいます。親しみやすいトーンで応答してください。",
    "デフォルト": "ユーザーの入力に対して、適切に応答してください。"
}

STYLE_INSTRUCTIONS = {
    "丁寧": "ユーザーは丁寧な言葉遣いを好みます。あなたも敬体（です・ます調）で応答してください。",
    "簡潔": "ユーザーは要点をまとめて話しています。あなたも簡潔に応答してください。",
    "創造的": "ユーザーは比喩や創造的な表現を使っています。あなたも表現を工夫して応答してください。",
    "ユーモラス": "ユーザーはユーモアを交えて話しています。あなたも遊び心のある応答をしてください。",
    "デフォルト": "ユーザーのスタイルに合わせて、自然に応答してください。"
}

def main():
    # --- 設定 ---
    llm_model_path = "./merged_model"
    sentence_model_name = "paraphrase-multilingual-MiniLM-L12-v2"
    classifier_name = "MoritzLaurer/mDeBERTa-v3-base-mnli-xnli"
    intent_labels = ["質問", "アイデアの要求", "感想", "雑談"]
    style_labels = ["丁寧", "簡潔", "創造的", "ユーモラス"]
    transition_log_file = "topic_transitions.jsonl"
    similarity_threshold = 0.4

    # --- モデルとツールの読み込み ---
    print("各モデルとツールを読み込んでいます...")
    model, tokenizer = load(llm_model_path)
    sentence_model = SentenceTransformer(sentence_model_name)
    classifier = pipeline("zero-shot-classification", model=classifier_name)
    janome_tokenizer = Tokenizer()
    print("モデルの読み込みが完了しました。")

    # --- 状態管理用変数 ---
    flashback_buffer = deque(maxlen=5)
    previous_response = None

    # --- チャットループ ---
    print("\nIlmチャットを開始します。終了するには 'exit' と入力してください。")
    while True:
        user_input = input("\nあなた: ")
        if user_input.lower() == 'exit':
            break

        prompt = ""
        flashback_triggered = False

        # --- フラッシュバックの検知 (最優先) ---
        for item in flashback_buffer:
            keyword = item['keyword']
            if keyword in user_input:
                print(f"(フラッシュバックを検知: {keyword})", end="")
                prompt = FLASHBACK_PROMPT.format(keyword=keyword, original_sentence=item['sentence'], user_input=user_input)
                flashback_triggered = True
                break

        # --- 通常のプロンプト組み立て ---
        if not flashback_triggered:
            # 意図とスタイルの検知
            intent_result = classifier(user_input, intent_labels, multi_label=False)
            style_result = classifier(user_input, style_labels, multi_label=False)
            detected_intent = intent_result['labels'][0]
            detected_style = style_result['labels'][0]
            print(f"(意図: {detected_intent} | スタイル: {detected_style})", end="")

            intent_instruction = f"\n### ユーザーの意図: {detected_intent}\n{INTENT_INSTRUCTIONS.get(detected_intent, INTENT_INSTRUCTIONS['デフォルト'])}"
            style_instruction = f"\n### ユーザーの対話スタイル: {detected_style}\n{STYLE_INSTRUCTIONS.get(detected_style, STYLE_INSTRUCTIONS['デフォルト'])}"
            transition_instruction = ""

            # 話題遷移の検知
            if previous_response:
                embedding_prev = sentence_model.encode(previous_response, convert_to_tensor=True)
                embedding_curr = sentence_model.encode(user_input, convert_to_tensor=True)
                cosine_sim = util.pytorch_cos_sim(embedding_prev, embedding_curr).item()
                print(f" (類似度: {cosine_sim:.2f})", end="")

                if cosine_sim < similarity_threshold:
                    print(" (話題の変化を検知)")
                    with open(transition_log_file, 'a', encoding='utf-8') as f:
                        f.write(json.dumps({"p": previous_response, "c": user_input, "s": cosine_sim}, ensure_ascii=False) + '\n')
                    transition_instruction = "\nスムーズな移行文を生成してから、ユーザーの質問に答えてください。\n" + TRANSITION_CONTEXT.format(previous_response=previous_response)
                else:
                    print()
            
            prompt = BASE_PROMPT.format(
                intent_instruction=intent_instruction,
                style_instruction=style_instruction,
                transition_instruction=transition_instruction,
                user_input=user_input
            )

        # --- 応答生成 ---
        print("\nIlm: ", end="", flush=True)
        current_response = ""
        for token in generate(model, tokenizer, prompt=prompt, verbose=False):
            current_response += token
            print(token, end="", flush=True)
        print()

        # --- 応答からキーワードを抽出し、バッファに保存 ---
        try:
            for token in janome_tokenizer.tokenize(current_response):
                if token.part_of_speech.startswith('名詞'):
                    flashback_buffer.append({'keyword': token.surface, 'sentence': current_response})
                    break
        except Exception:
            pass

        previous_response = current_response

if __name__ == "__main__":
    main()