Create app.py

app.py

@@ -0,0 +1,270 @@
+import os
+import shutil
+import torch
+from transformers import pipeline, AutoTokenizer, BitsAndBytesConfig
+import gradio as gr
+import logging
+from pathlib import Path
+import json
+from typing import Tuple # generate_response の型ヒントのために追加
+
+# ragsys03.py から RAGSystem クラスと QueryResult をインポート
+from ragsys03 import RAGSystem, QueryResult 
+
+# ロギング設定
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# グローバル変数
+rag_system: RAGSystem = None
+llm_pipeline = None
+
+# RAGデータ保存用のベースディレクトリ
+RAG_BASE_DIR = Path('rag_data')
+
+def initialize_systems():
+    """
+    RAGシステムとLLMパイプラインを初期化します。
+    """
+    global rag_system, llm_pipeline
+
+    if rag_system is None:
+        logger.info("RAGSystemを初期化中...")
+        try:
+            # 最新のインデックスをロードするか、新しいディレクトリを作成する
+            rag_system = RAGSystem(
+                model_name='all-MiniLM-L6-v2', # SentenceTransformer モデル
+                index_type='ivf', # または 'flat'
+                n_clusters=100, # IVFの場合のクラスタ数
+                index_base_dir=RAG_BASE_DIR,
+                load_latest=True # 最新のインデックスを自動的にロード
+            )
+            logger.info("RAGSystemの初期化が完了しました。")
+            
+            # 初期化時にインデックスが存在しない場合、ここで構築を促す
+            if rag_system.index is None or rag_system.index.ntotal == 0:
+                logger.warning("RAGシステムにインデックスがありません。文書をアップロードして「インデックスを構築」ボタンを押してください。")
+
+        except Exception as e:
+            logger.critical(f"RAGSystemの初期化中にエラーが発生しました: {e}")
+            rag_system = None # 初期化失敗時はNoneに設定
+            # Gradio UIでエラーメッセージを表示するための処理を考慮するか、
+            # 各アクションでrag_systemがNoneの場合の処理を堅牢にする
+            # ここでは単にログに出力し、後続の関数でNoneチェックを行う
+
+    if llm_pipeline is None:
+        logger.info("LLMパイプラインを初期化中: rinna/japanese-gpt-neox-3.6b-instruction-sft")
+        try:
+            # 量子化を無効にした設定
+            llm_pipeline = pipeline(
+                "text-generation",
+                model="rinna/japanese-gpt-neox-3.6b-instruction-sft",
+                tokenizer=AutoTokenizer.from_pretrained("rinna/japanese-gpt-neox-3.6b-instruction-sft", use_fast=False),
+                torch_dtype=torch.bfloat16, # CPUでも動くようにbfloat16のまま
+                device_map="auto", # autoのままでOK (GPUがなければCPUに自動的に割り当たる)
+                # model_kwargs={"quantization_config": quantization_config} を削除！
+            )
+            logger.info("LLMパイプラインの初期化が完了しました。")
+        except Exception as e:
+            logger.critical(f"LLMパイプラインの初期化中にエラーが発生しました: {e}")
+            llm_pipeline = None # 初期化失敗時はNoneに設定
+            # 同様に、エラー発生時はllm_pipelineをNoneのままにする
+
+def generate_response(query: str, top_k: int, similarity_threshold: float) -> Tuple[str, str]:
+    """
+    RAGシステムとLLMを使用して応答を生成します。
+    Args:
+        query (str): ユーザーからのクエリ。
+        top_k (int): 検索する文書の数。
+        similarity_threshold (float): コサイン類似度の閾値。
+    Returns:
+        Tuple[str, str]: LLMの回答と、検索された文書のテキスト。
+    """
+    if rag_system is None or llm_pipeline is None:
+        return "システムが初期化されていません。アプリケーションを再起動してください。", ""
+
+    logger.info(f"クエリ処理開始: {query}")
+    query_result: QueryResult = rag_system.query(query, top_k=top_k, similarity_threshold=similarity_threshold)
+
+    llm_prompt = query_result['llm_prompt']
+    retrieved_docs = query_result['retrieved_documents']
+    search_time = query_result['search_time']
+
+    if not retrieved_docs:
+        llm_answer = "関連する情報が見つかりませんでした。別の表現で質問していただくか、より具体的な内容で質問してください。"
+        return llm_answer, "関連文書なし"
+
+    try:
+        outputs = llm_pipeline(
+            llm_prompt,
+            max_new_tokens=512,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
+            repetition_penalty=1.1,
+            num_return_sequences=1,
+            pad_token_id=llm_pipeline.tokenizer.eos_token_id # pad_token_idを指定
+        )
+
+        llm_answer = outputs[0]['generated_text']
+        # プロンプト部分を除去して回答のみを抽出 (モデルの出力形式による)
+        if llm_answer.startswith(llm_prompt):
+            llm_answer = llm_answer[len(llm_prompt):].strip()
+        
+        logger.info(f"LLM応答生成完了 (検索時間: {search_time:.4f}秒)")
+        
+        retrieved_docs_text_formatted = "\n\n".join([f"**文書{i+1}:** {doc}" for i, doc in enumerate(retrieved_docs)])
+        
+        return llm_answer, retrieved_docs_text_formatted
+
+    except Exception as e:
+        logger.error(f"LLMによる応答生成中にエラーが発生しました: {e}")
+        return f"応答生成中にエラーが発生しました: {str(e)}", ""
+
+def upload_documents(file):
+    """
+    アップロードされたJSONファイルから文書をロードし、インデックスを再構築します。
+    """
+    global rag_system
+
+    if rag_system is None:
+        # RAGSystemが未初期化の場合、ここで初期化を試みる
+        # initialize_systems() は、Gradioの起動時に一度しか呼ばれないため、
+        # ここでrag_systemがNoneの場合、初期化に失敗している可能性がある。
+        # 代わりに、エラーメッセージを返す
+        return "エラー: RAGシステムが初期化されていません。アプリケーションを再起動してください。", ""
+
+    if file is None:
+        return "エラー: ファイルがアップロードされていません。", ""
+
+    uploaded_file_path = Path(file.name)
+    logger.info(f"アップロードされたファイル: {uploaded_file_path}")
+
+    try:
+        # ragsys03.py の load_documents メソッドを使用
+        rag_system.load_documents(uploaded_file_path)
+        
+        # インデックス構築のメッセージを表示し、ユーザーにボタンを押してもらう
+        return (
+            f"'{uploaded_file_path.name}' から {len(rag_system.documents)} 件の文書を正常にロードしました。\n"
+            "続けて「インデックスを構築」ボタンを押してください。",
+            "文書がロードされました。インデックス構築が必要です。"
+        )
+    except Exception as e:
+        logger.error(f"文書のロード中にエラーが発生しました: {e}")
+        return f"文書のロード中にエラーが発生しました: {str(e)}", ""
+
+def build_index_action():
+    """
+    RAGインデックスを構築します。
+    """
+    global rag_system
+    if rag_system is None:
+        return "エラー: RAGシステムが初期化されていません。", ""
+
+    if not rag_system.documents:
+        return "エラー: インデックス作成用の文書がロードされていません。", ""
+
+    try:
+        # force_rebuild=True で常に再構築
+        rag_system.build_index(force_rebuild=True) 
+        stats = rag_system.get_stats()
+        return (
+            f"インデックスが正常に構築されました。\n"
+            f"総文書数: {stats.get('documents_count', 'N/A')}\n"
+            f"インデックスタイプ: {stats.get('index_type', 'N/A')}\n"
+            f"インデックスディレクトリ: {stats.get('current_index_dir', 'N/A')}\n"
+            f"ファイルサイズ: {stats.get('index_file_size_mb', 'N/A')} MB"
+        ), "インデックス構築完了"
+    except Exception as e:
+        logger.error(f"インデックスの構築中にエラーが発生しました: {e}")
+        return f"エラー: インデックスの構築中に問題が発生しました - {str(e)}", ""
+
+def get_rag_stats_action():
+    """RAGシステムの統計情報を表示します。"""
+    global rag_system
+    if rag_system is None:
+        return "RAGシステムはまだ初期化されていません。"
+    try:
+        stats = rag_system.get_stats()
+        stats_str = "\n".join([f"{k}: {v}" for k, v in stats.items()])
+        return f"RAGシステムの状態:\n{stats_str}"
+    except Exception as e:
+        logger.error(f"RAG統計情報の取得中にエラーが発生しました: {e}")
+        return f"統計情報の取得中にエラーが発生しました: {str(e)}"
+
+# Gradio UI の構築
+with gr.Blocks() as demo:
+    gr.Markdown(
+        """
+        # RAG (Retrieval Augmented Generation) デモ
+        rinna/japanese-gpt-neox-3.6b-instruction-sft と FAISS を利用したRAGシステムです。
+        JSONファイルをアップロードして独自の知識ベースを構築し、質問をすることができます。
+
+        ## 使用方法
+        1. 「文書をアップロード」セクションで、`documents` キーに文字列のリストを持つJSONファイルをアップロードします。
+        2. 「インデックスを構築」ボタンをクリックして、アップロードされた文書から検索インデックスを作成します。
+        3. 「RAG質問」セクションで質問を入力し、「質問を送信」ボタンをクリックします。
+        """
+    )
+
+    with gr.Tab("RAG質問"):
+        gr.Markdown("### RAG質問")
+        with gr.Row():
+            query_input = gr.Textbox(label="質問を入力してください", placeholder="RAGシステムの主な利点は何ですか？", lines=2)
+            submit_button = gr.Button("質問を送信")
+        
+        with gr.Row():
+            top_k_slider = gr.Slider(minimum=1, maximum=10, value=3, step=1, label="取得文書数 (top_k)")
+            similarity_threshold_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, step=0.05, label="類似度閾値 (0.0-1.0)")
+
+        llm_output = gr.Textbox(label="LLMの回答", lines=10, interactive=False)
+        retrieved_docs_output = gr.Markdown("検索された関連文書", visible=True)
+
+    with gr.Tab("文書管理"):
+        gr.Markdown("### 文書とインデックス管理")
+        with gr.Row():
+            file_upload_button = gr.File(label="文書JSONファイルをアップロード", file_types=[".json"])
+            upload_status = gr.Textbox(label="アップロードステータス", interactive=False)
+        
+        with gr.Row():
+            build_index_button = gr.Button("インデックスを構築")
+            build_index_status = gr.Textbox(label="インデックス構築ステータス", interactive=False)
+        
+        gr.Markdown("### RAGシステム情報")
+        get_stats_button = gr.Button("RAGシステムの状態を表示")
+        rag_stats_output = gr.Textbox(label="RAGシステム情報", interactive=False, lines=5)
+
+
+    # イベントハンドラの登録
+    submit_button.click(
+        fn=generate_response,
+        inputs=[query_input, top_k_slider, similarity_threshold_slider],
+        outputs=[llm_output, retrieved_docs_output]
+    )
+
+    file_upload_button.upload(
+        fn=upload_documents,
+        inputs=file_upload_button,
+        outputs=[upload_status, retrieved_docs_output] # アップロード結果と関連文書表示エリアを更新
+    )
+
+    build_index_button.click(
+        fn=build_index_action,
+        inputs=[],
+        outputs=[build_index_status, retrieved_docs_output] # 構築結果と関連文書表示エリアを更新
+    )
+
+    get_stats_button.click(
+        fn=get_rag_stats_action,
+        inputs=[],
+        outputs=rag_stats_output
+    )
+
+# システムの初期化をアプリケーション起動時に実行
+# initialize_systems() の中でエラーが発生した場合、rag_system や llm_pipeline が None になる
+# これにより、後続の関数呼び出しでこれらの変数がNoneであることのチェックが機能する
+initialize_systems()
+
+if __name__ == "__main__":
+    demo.launch(debug=True, share=True)