wikipedia-japanese-open-rag

Wikipediaの日本語記事を元に、ユーザの質問に回答するGradioベースのRAGのサンプル。

実装の背景と詳細の解説

コーパス

• fivehints氏によるsingletongue/wikipedia-utils
• hotchpotch氏による同データのembeddings,faiss index

使用モデル

• intfloat/multilingual-e5-large
• elyza/ELYZA-japanese-Llama-2-13b-instruct

使い方

環境

torch.cuda.is_available()かつ、VRAM 96GB程度のマシンでの稼働を想定している。
ただし、使用するVRAMの大部分はvllmが利用するcache領域であるため、gpu_memory_utilizationを指定することでVRAM 40GB程度まで必要量を削減可能。(ELYZA-japanese-Llama-2-13b-instructが30GB弱、multilingual-e5-largeが10GB弱)

        engine_args = AsyncEngineArgs(model=LLM_MODEL_ID, dtype='bfloat16',
                                      tensor_parallel_size=4,
                                      disable_log_requests=True,
                                      disable_log_stats=True,
                                      gpu_memory_utilization=0.6) # ←ここでVRAM全体の何％を推論に割り当てるかを指定する(ここでは60%)

GPU VRAMは必ずしも単一のGPUで供給する必要はなく、マシンが複数のGPUを利用可能な場合はtensor_parallel_sizeを指定することで分散的に推論できる。

        engine_args = AsyncEngineArgs(model=LLM_MODEL_ID, dtype='bfloat16',
                                      tensor_parallel_size=4, # ←ここでGPUの数を指定
                                      disable_log_requests=True,
                                      disable_log_stats=True,
                                      gpu_memory_utilization=0.6)

Stand AloneモードとAPIモード

GradioのUIと推論エンジンを同一のPythonプロセスで稼働させるStand Aloneモードと、UIと推論を異なるプロセスに分割して連携させるAPIモードを用意している。Stand Aloneモードはシンプルに運用可能な反面、スケーラビリティに限界がある。
APIモードでは、推論部分をFast APIに切り出しており、n台のノードにAPIをデプロイし、適切にロードバランスしてGradioのUIからcallすることでスケールアウトするような用途を想定している。

Stand Aloneモード

pip install -r requirements.txt
python app.py # 同一プロセスでGUIとRAG(検索、推論)を実行

APIモード

pip install -r requirements.txt
python app.py --api_mode # UIのみを起動。検索、推論時はgunicornのapiを呼び出す
gunicorn -w 1 --timeout 0 -k uvicorn.workers.UvicornWorker rag_inf_api:app # 適宜別ノードから起動