RAG向けデータベースを作成・評価

# JupyterLabが使えるDockerイメージ
# 自分がこれまで作成したNoteBookに必要なパッケージをインストールしたもの。

FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04

# Set bash as the default shell
ENV SHELL=/bin/bash

# Build with some basic utilities
RUN apt update \
	&& apt install -y \
        wget \
        bzip2 \
        git \
        git-lfs \
        curl \
        unzip \
        file \
        xz-utils \
        sudo \
        python3 \
        python3-pip && \
        apt-get autoremove -y && \
        apt-get clean && \
        rm -rf /usr/local/src/*

# alias python='python3'
RUN ln -s /usr/bin/python3 /usr/bin/python

RUN pip install --upgrade pip setuptools \
 	&& pip install torch==2.2.2 torchvision==0.17.2 torchaudio==2.2.2 \
	--index-url https://download.pytorch.org/whl/cu121 \
	&& pip install torch torchvision torchaudio \
	&& pip install jupyterlab matplotlib pandas scikit-learn ipywidgets \
	&& pip install transformers accelerate sentencepiece einops \
	&& pip install langchain bitsandbytes protobuf \
	&& pip install auto-gptq optimum \
	&& pip install pypdf tiktoken sentence_transformers faiss-gpu trafilatura \
	&& pip install langchain-community langchain_openai wikipedia \
	&& pip install langchain-huggingface unstructured html2text rank-bm25 janome \
	&& pip install langchain-chroma sudachipy sudachidict_full \
	&& pip install mysql-connector-python

# Install llama-cpp-python[server] with cuBLAS on
#RUN CMAKE_ARGS="-DLLAMA_CUBLAS=on" FORCE_CMAKE=1 \
#        pip install llama-cpp-python[server]==0.2.75 --force-reinstall --no-cache-dir
RUN CMAKE_ARGS="-DGGML_CUDA=on" FORCE_CMAKE=1 \
        pip install llama-cpp-python[server] --force-reinstall --no-cache-dir

RUN pip install -U numpy==1.26.4 

# Create a working directory
WORKDIR /workdir

# Port number in container side
EXPOSE 8888

ENTRYPOINT ["jupyter-lab", "--ip=0.0.0.0", "--port=8888", "--no-browser", "--allow-root", "--NotebookApp.token=''"]

CMD ["--notebook-dir=/workdir"]

# Vector databaseを作成する。

import time

start_time = time.time()

#　Wikipediaから天文関係の記事のみを抽出したテキストデータは、WIKI_TEXTとして事前に準備。
# VectorStoreするため、テキストコーパスをチャンクに分割。
WIKI_TEXT = "wiki_text"

# チャンク文字数、オーバラップ文字数を定義
CHUNK_SIZE = 300
CHUNK_OVERLAP = 30

from langchain.text_splitter import RecursiveCharacterTextSplitter

with open(WIKI_TEXT,'r',encoding='utf-8') as f:
    text = f.read()

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size = CHUNK_SIZE,  # チャンクの文字数
    chunk_overlap = CHUNK_OVERLAP,  # チャンクオーバーラップの文字数
)

texts = text_splitter.split_text(text)

WIKI_TEXT = "wiki_text"
WIKI_DB = "wiki_vs.db"

## データベースを構築

from langchain_huggingface import HuggingFaceEmbeddings

model_name = "intfloat/multilingual-e5-large"
model_path = f"/workdir/models/{model_name}"

embeddings = HuggingFaceEmbeddings(
    model_name = model_path,
    model_kwargs = {'device':'cuda:0'},
)

from langchain_community.vectorstores import FAISS

# ベクトルデータベースを構築し、WIKI_DBとして保存する

db = FAISS.from_texts(texts, embeddings)

end_time = time.time()

db.save_local(WIKI_DB)

processing_time = end_time - start_time

print("processing_time(sec): ", processing_time)

processing_time(sec):  704.8360137939453

# キーワード検索用のデータベースを作成
#
import time

start_time = time.time()

# my_preprocess_funcを定義

from janome.tokenizer import Tokenizer

t = Tokenizer()

def my_preprocess_func(text):
    keywords = []
    for token in t.tokenize(text):
        pos = token.part_of_speech.split(',')[0]
        if (pos in ["名詞", "動詞", "形容詞"]):
            keywords.append(token.surface)
    return keywords

#　Wikipediaから天文関係の記事のみを抽出したテキストデータは、WIKI_TEXTとして事前に準備。
# キーワードベースのデータベースを作成するため、テキストコーパスをチャンクに分割。
WIKI_TEXT = "wiki_text"

# チャンク文字数、オーバラップ文字数を定義
CHUNK_SIZE = 300
CHUNK_OVERLAP = 30

from langchain.text_splitter import RecursiveCharacterTextSplitter

with open(WIKI_TEXT, encoding='utf-8') as f:
    text = f.read()

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size = CHUNK_SIZE,  # チャンクの文字数
    chunk_overlap = CHUNK_OVERLAP,  # チャンクオーバーラップの文字数
)

texts = text_splitter.split_text(text)

from langchain_community.retrievers import BM25Retriever

WIKI_DB = "wiki_kw.pkl"

db = BM25Retriever.from_texts(
    texts,
    preprocess_func=my_preprocess_func,
)

end_time = time.time()

import pickle
with open(WIKI_DB, 'wb') as f:
    pickle.dump(db, f)

processing_time = end_time - start_time

print("processing_time(sec): ", processing_time)

processing_time(sec):  753.1111702919006

# ベクトルデータベース検索結果を評価する
# 

WIKI_DB = "wiki_vs.db"
import os
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS

model_name = "intfloat/multilingual-e5-large"
model_path = f"/workdir/models/{model_name}"

embeddings = HuggingFaceEmbeddings(
    model_name = model_path,
    model_kwargs = {'device':'cuda:0'},
)

# 事前に構築したベクトルデータベースを読み込む
if os.path.exists(WIKI_DB):
    db = FAISS.load_local(WIKI_DB, embeddings, allow_dangerous_deserialization=True)
else:
    print("="*80)
    print("You need to make vector database")
    print("="*80)

# ベクトル検索で4文書を得る
query = "r過程について教えてください"

docs = db.similarity_search(query, k=4)
print(f"1st: {docs[0].page_content}")
print(f"2nd: {docs[1].page_content}")
print(f"3rd: {docs[2].page_content}")
print(f"4th: {docs[3].page_content}")

# ベクトル検索で4文書を得る
query = "ナンシー・グレース・ローマン宇宙望遠鏡について教えてください"

docs = db.similarity_search(query, k=4)
print(f"1st: {docs[0].page_content}")
print(f"2nd: {docs[1].page_content}")
print(f"3rd: {docs[2].page_content}")
print(f"4th: {docs[3].page_content}")

# ベクトル検索で4文書を得る
query = "B2FH論文について教えてください"

docs = db.similarity_search(query, k=4)
print(f"1st: {docs[0].page_content}")
print(f"2nd: {docs[1].page_content}")
print(f"3rd: {docs[2].page_content}")
print(f"4th: {docs[3].page_content}")

# キーワード検索結果を評価する
# 

from janome.tokenizer import Tokenizer

t = Tokenizer()

def my_preprocess_func(text):
    keywords = []
    for token in t.tokenize(text):
        pos = token.part_of_speech.split(',')[0]
        if (pos in ["名詞", "動詞", "形容詞"]):
            keywords.append(token.surface)
    return keywords

WIKI_DB = "wiki_kw.pkl"
import os
import pickle
from langchain_community.retrievers import BM25Retriever

# 事前に構築したキーワードのデータベースを読み込む
if os.path.exists(WIKI_DB):
    with open(WIKI_DB, 'rb') as f:
        retriever = pickle.load(f)
else:
    print("="*80)
    print("You need to make keyword database")
    print("="*80)

# キーワード検索で4文書を得る
query = "r過程について教えてください"

docs = retriever.get_relevant_documents(query, k=4)
print(f"1st: {docs[0].page_content}")
print(f"2nd: {docs[1].page_content}")
print(f"3rd: {docs[2].page_content}")
print(f"4th: {docs[3].page_content}")

# キーワード検索で4文書を得る
query = "ナンシー・グレース・ローマン宇宙望遠鏡について教えてください"

docs = retriever.get_relevant_documents(query, k=4)
print(f"1st: {docs[0].page_content}")
print(f"2nd: {docs[1].page_content}")
print(f"3rd: {docs[2].page_content}")
print(f"4th: {docs[3].page_content}")

# キーワード検索で4文書を得る
query = "B2FH論文について教えてください"

docs = retriever.get_relevant_documents(query, k=4)
print(f"1st: {docs[0].page_content}")
print(f"2nd: {docs[1].page_content}")
print(f"3rd: {docs[2].page_content}")
print(f"4th: {docs[3].page_content}")