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)