はじめに
今回は最近、流行りのRAG(Retrieval-Augmented Generation)をLangChainなどを使用せずに処理を構築した内容を紹介する実装編になります。前回は理解編としてRAGの処理の流れを紹介しました。良ければ前回の記事と合わせてみて頂けると幸いです。
前回の記事: RAGをLangChainを使わずに作ってみる（理解編）
RAGを実装する
処理の概要
RAG の処理は以下の 2 つの処理に分けることができ、それぞれの大まかな処理は以下でした。
-
Inferencing: RAGをつかったLLMの推論
- LLMモデルに質問(prompt)の投げる
- 質問(prompt)を埋め込みモデルでベクトルに変換する
- 変換したベクトルをもとにDB内を検索する
- 検索結果の引用情報を質問(prompt)に埋め込む
- 引用情報を埋め込んだ質問(prompt)を用いて推論を行う
-
Ingestion: ベクトルDBの準備
- 関連文書を任意のサイズ(chunk_size)で分割する
- 関連文書を埋め込みベクトルに変換する
- DBに関連文書とベクトルを記録する
処理を図示すると以下のように表現でき、赤く囲った部分がIngestion処理、青く囲った部分がInferencing処理となります。
今回はこれらの処理を実装する主なライブラリしてモデルまわりはhugging face関連のものを利用し。DBはelastic searchを用い、dockerを用いて環境の構築を行います。
使用する環境について
ベクトルDBに関して
elastic searchコンテナの立ち上げ
ベクトル DB は elastic search を使用し関連文書の検索、ベクトルの管理を行います。
DBは以下のdockerを用いて環境を構築します。docker-composeではデータの確認用としてkibanaも合わせて起動させています。
elasticsearchが起動後にhttp://localhost:5601/app/home#/へアクセスしてデータの確認ができると思います。
dockerfile
FROM docker.elastic.co/elasticsearch/elasticsearch:8.1.0
RUN elasticsearch-plugin install analysis-kuromoji
docker-compose.yml
version: '3'
services:
elasticsearch:
container_name: elasticdb
build:
context: ./elasticsearch
dockerfile: dockerfile
environment:
- node.name=es01
- cluster.name=es-docker-cluster
- discovery.type=single-node
- xpack.security.enabled=false
ports:
- 9200:9200
ulimits:
memlock:
soft: -1
hard: -1
volumes:
- ./elasticsearch/esdata01:/usr/share/elasticsearch/data
kibana:
image: docker.elastic.co/kibana/kibana:8.1.0
ports:
- 5601:5601
environment:
ELASTICSEARCH_HOSTS: http://elasticdb:9200
ディレクトリ構成
.
├── /elastchsearch
│ ├── /esdata01
│ └── dockerfile
└── docker-compose.yml
elasticsearchでのベクトルの扱い方
elasticsearchではベクトル検索の機能があり、その機能を利用するための保存形式があります。
ここで必要になる情報として、埋め込みモデルが埋め込みベクトルを生成する次元数がいくつになるかが必要になります。
elasticsearchでベクトル情報を扱うには、保存形式の"type"を"dense_vector"に設定することに加えて、"dims"に埋め込みベクトルの次元数を設定する必要があります。
elasticsearchに登録するindex（文書）の保存情報の定義は今回は以下のようにしました。
mapping_info = {
"mappings": {
"properties": {
"file_title":{
"type": "text"
},
"context":{
"type": "text"
},
"context_vector":{
"type": "dense_vector",
"dims": 384
}
}
}
}
Pythonでindexを登録する処理は以下になります。
from elasticsearch import Elasticsearch
es = Elasticsearch('http://localhost:9200')
es.indices.create(index="sample_index", body=mapping_info)
モデルに関して
埋め込みモデルと推論モデルは hugging face からモデルを参照して使用します。
今回は一連の処理を実装することを目的としてるので、モデルの精度は考慮しません。
モデルは以下を使用しました。
- 埋め込みモデル: sentence-transformers/all-MiniLM-L6-v2
- 推論モデル：cyberagent/open-calm-small
Ingestion
まずは既存文書を埋め込みベクトルへ変換し、DBに保存する処理について紹介します。Ingestionで必要な処理は以下でした。
- Ingestion: ベクトルDBの準備
- 関連文書を任意のサイズ(chunk_size)で分割する
- 関連文書を埋め込みベクトルに変換する
- DBに関連文書とベクトルを記録する
ここではVectorDBクラスを定義し、処理を実装していきます。
なお、細かいDBでのデータ管理は考慮していないので、文書のベクトルを扱う最低限の処理の実装の紹介になります。
1. 関連文書を任意のサイズ(chunk_size)で分割する
ここではベクトルに変換する既存文書を任意のサイズ(chunk_size)で分割する処理になります。
chunk_sizeで分割する意図としては検索する情報量を調整するためになります。検索結果として文書全体が返ってきたとしても、具体的にどの部分を引用すれば良いかがわからないと思います。なので、検索結果として適切な粒度となるように分割し、埋め込みベクトルに変化します。
chunk_sizeとしては文書のページ単位や段落単位など様々な粒度が考えられます。今回は文字数で情報を区切ることとしました。
CHUNK_SIZE =256 # 元のデータを分割する文字
class VectorDB:
def __init__(self):
self.chunk_size = CHUNK_SIZE
def split_chunk(self, document):
self.chunk_size
sentences = []
for doc in document:
doc_length = len(doc)
start_idx = 0
end_idx = self.chunk_size
while doc_length > 0:
sentences.append(doc[start_idx:end_idx])
doc_length -= self.chunk_size
start_idx = end_idx
if doc_length < self.chunk_size:
end_idx = -1
else:
end_idx += self.chunk_size
return sentences
2. 関連文書を埋め込みベクトルに変換する
今回はVectorDBクラスで埋め込みモデルを管理するように実装します。なので、ここで必要な行程は以下の2つになります。
- 埋め込みモデルの読み込み
- 関連文書のベクトル化
埋め込みモデルの読み込み
ここではhuggingfaceのtransformesライブラリを用いてモデルの読み込み処理を行います。
VectorDBクラスでmodelをロードし、管理するような実装とします。
処理としては単純で、load_embedding_model()でhugging faceから埋め込みモデルを読み込むsentence_transformersを使用して読み込んだモデルを返す関数を実装しています。
必要なコードとしては以下の内容になります。
from sentence_transformers import SentenceTransformer
EMBEDDED_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
class VectorDB:
def __init__(self):
self.embedding_model = self.load_embedding_model(EMBEDDED_MODEL_NAME)
def load_embedding_model(self, model_name: str):
print(f">> load embedding model from huggingface: {model_name}")
return SentenceTransformer(model_name)
関連文書のベクトル化
ここでは先ほどの1.でchunkingした関連文書を受け取り、埋め込みベクトルに変換する部分になります。
コードの内容としては単純で、関連文書を引数にencodeメソッドを呼ぶだけになります。
class VectorDB:
def conv2vec(self, sentences: str):
embeddings = self.embedding_model.encode(sentences)
return embeddings
3. DBに関連文書とベクトルを記録する
ここではelastic searchまわりの処理を見ていきます。elastic searchのindexの設定などは先ほどのベクトルDBに関してと同じ内容を使用します。
- elastic searchへの接続
- 埋め込みベクトルをDBへ記録
elastic searchへの接続
ここではelastic searchないのindexを操作するインスタンスを作成します。BDへの接続はconect_db()内で実装しています。また、DB内に今回使用するindexが存在しない場合に、作成する処理を組み込んでいます。
from elasticsearch import Elasticsearch
DB_URI = 'http://localhost:9200'
INDEX_NAME = "sample_index"
INDEX_DATA = "mapping.json"
class VectorDB:
def __init__(self):
self.conect_db(DB_URI)
self.index_name = INDEX_NAME
# indexがなければ作成する
if not self.es.indices.exists(index=self.index_name):
self.es.indices.create(index=self.index_name, body=INDEX_DATA)
def conect_db(self, uri):
self.es = Elasticsearch(uri)
print(">> Connect to elastic search")
埋め込みベクトルをDBへ記録
ここではelastic searchのindexの項目に合わせてデータを記録する処理になります。
register_document_vector()の処理としては、titleで文書のタイトルとdocumentで文書内の文字列をそれぞれ引数として受け取っています。
受け取ったdocumentを先ほどのchunking処理と、埋め込みベクトルの変換処理をはじめの2行で行っています。そして、elastic searchのindexの項目に合わせて、分割されたデータごとにdictを作成します。ここは分割したデータが識別できるようにtitleを設定するようにしています。
なので、文書のタイトル、chunk_sizeされた元の文、ベクトル化された文の3つの情報のdictが1つのデータとしてDBに記録されます。
def register_document_vector(self, title: str, document: list):
sentences = self.split_chunk(document)
vector = self.conv2vec(sentences)
index_data = []
for i in tqdm(range(len(sentences)), desc=f"[registration / title:{title}]"):
index_data.append(
{
"file_title": title,
"context": sentences[i],
"context_vector": vector[i],
}
)
self.register(index_data)
register()は先ほどのdictをelastic searchに記録する処理です。記録方法はelastic searchのhelpersを用いて複数のデータを記録する処理になっています。記録方法の詳細に関してはココを参考にしました。
from elasticsearch import helpers
class VectorDB:
def register(self, indexes: list):
def gendata(indexe_data: list):
for data in indexe_data:
yield {
"_ope_type": "index",
"_index": self.index_name,
"_source": data,
}
helpers.bulk(self.es, gendata(indexes))
Ingestionの実行
完成したVectorDBクラスのregister_document_vector()を対して以下のように呼び出すことで、文書単位でDBに情報を登録できます。
from data import preporcess as preprc
from database.vector_db import VectorDB
def ingestion():
pptx_data_list = preprc.preporcess() # 複数の文書に対して文字を抽出する処理
db = VectorDB()
for pptx_data in pptx_data_list: # 1つずつ文書を取りだして登録する
db.register_document_vector(
title=pptx_data["file_name"], document=pptx_data["text_list"]
)
if __name__ == "__main__":
ingestion()
Inferencing
ここでは引用情報を検索し、推論処理を行うInferencingにあたる部分について紹介します。Inferencingの処理の流れは以下の内容が必要でした。
- Inferencing: RAGをつかったLLMの推論
- LLMモデルに質問(prompt)の投げる
- 質問(prompt)を埋め込みモデルでベクトルに変換する
- 変換したベクトルをもとにDB内を検索する
- 検索結果の引用情報を質問(prompt)に埋め込む
- 引用情報を埋め込んだ質問(prompt)を用いて推論を行う
ここでの実装の方針として、DBに関連する処理は先ほどのVectorDBクラスを拡張し、モデルによる推論部分はInferencingModelクラスに担ってもらうようにします。これらの処理を一つにあわせた処理をInferenceLogicクラスとして、すべての処理を包括して実行できるように実装していきたいと思います。
VectorDBクラスを拡張する処理は2.と3.の処理にあたり、InferencingModelクラスには4.と5.の処理を実装します。また、1.に関してはユーザー側の操作の処理になるので今回は実装しません。質問(prompt)の入力に関してはInferenceLogicのインスタンスで文字列を受け取ることを想定してすすめます。
2. 質問(prompt)を埋め込みモデルでベクトルに変換する
ここで必要になる処理は以下の2つになります。
- 埋め込みモデルの読み込み
- 質問(prompt)の埋め込みベクトルへの変換
これらの処理は既にIngestionの処理で実装済みになっています。埋め込みモデルはVectorDBクラスで利用できるようになっており、質問(prompt)の埋め込みベクトルへの変換も関連文書のベクトル化と同様の機能を利用すれば問題ありません。
3. 変換したベクトルをもとにDB内を検索する
ここではelastic searchの機能を用いてベクトル検索を行います。elastic searchへのクエリは以下のようなdictに要件を記述します。"params": {"query_vector": input_vec}のinput_vecが質問(prompt)の埋め込みベクトルにあたり、"size": self.contex_numが類似度の高い引用情報を上からいくつ取得するかを設定しています。
- elastic searchへのクエリ
{ "query": { "script_score": { "query": {"match_all": {}}, "script": { "source": "cosineSimilarity(params.query_vector, 'context_vector') + 1.0", "params": {"query_vector": input_vec}, # 検索クエリのベクトル }, } }, "size": self.contex_num, # 取得する件数 }
elastic searchへのクエリを組み込んだ、検索処理が以下のsearch_similary_sentences()になります。引数のinput_dataは質問(prompt)の文字列になります。また、検索処理はself.es.search(index=self.index_name, body=query) で行っており、indexと先ほどのクエリを指定して検索を行います。検索結果から引用する情報の文字列部分だけをcontext_textに取り出して、取得件数分のリストにしてreturnしています。
class VectorDB:
def __init__(self):
self.contex_num = CONTEXT_NUM
def search_similary_sentences(self, input_data):
input_vec = self.conv2vec(input_data)
query = {
"query": {
"script_score": {
"query": {"match_all": {}},
"script": {
"source": "cosineSimilarity(params.query_vector, 'context_vector') + 1.0",
"params": {"query_vector": input_vec}, # 検索クエリのベクトル
},
}
},
"size": self.contex_num, # 取得する件数
}
results = self.es.search(index=self.index_name, body=query) # elasticsearchでの検索
context_text = []
for hits in results["hits"]["hits"]:
context_text.append(hits["_source"]["context"])
return context_text
4. 検索結果の引用情報を質問(prompt)に埋め込む
ここからは推論モデルを扱うInferencingModelクラスを実装していきます。hugging faceからロードして使うモデルに合わせてinitの処理や推論方法に関しては適宜変えてください。今回は一例として処理を紹介します。
promptの処理ではあらかじめRAG用に独自に定義しておいたprompt_templateに対して、formatを利用してユーザーの質問とDBから検索した引用情報を埋め込んでいます。
また、モデルに応じて入力されるpromptの形式があらかじめ指定されている場合がるので、その内容に従って独自に定義しておいたprompt_templateの部分がユーザーからの入力となるように処理しています。
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
INFERENCE_MODEL_NAME = "cyberagent/open-calm-small"
PROMPT_TEMPLATE = """
## 指示: \n以下の「引用情報」を元に質問に回答してください。なお、「引用情報」に無い情報は回答に含めないでください。\n
## 引用情報: \n{context}\n
## 質問: \n{question} \n
"""
class InferencingModel:
def __init__(self):
self.set_prompt_template(PROMPT_TEMPLATE)
def set_prompt_template(self, template):
self.prompt_template = template
return self.prompt_template
def generate_prompt(self, question, context_list):
context = ""
for i in range(len(context_list)):
context += context_list[i] + "\n"
_input_prompt = [
{
"speaker": "ユーザー",
"text": self.prompt_template.format(question=question, context=context),
}
]
input_prompt = [f"{uttr['speaker']}: {uttr['text']}" for uttr in _input_prompt]
input_prompt = self.nl_label.join(input_prompt)
input_prompt = input_prompt + self.nl_label + self.model_label
return input_prompt
5. 引用情報を埋め込んだ質問(prompt)を用いて推論を行う
推論用のモデルの読み込みと推論処理になります。推論用のモデルの読み込みはload_model()でtokenizerとモデル本体を読み込んでいます。
推論部分はinference()になり、引数のquestionはユーザーの質問、contextはVectorDBクラスのsearch_similary_sentences()の戻り値を受け取り推論結果を返すようになっています。inference()の1行目で先ほどの引用情報を埋め込んだ質問(prompt)が生成され、その情報をもと推論が行われてoutputとして値が返却されるようになています。
この推論処理はhugging faceのモデルカードに記載されているので、その部分を参考に実装してみてください。
mport torch
from transformers import AutoModelForCausalLM, AutoTokenizer
INFERENCE_MODEL_NAME = "cyberagent/open-calm-small"
class InferencingModel:
def __init__(self):
self.load_model(INFERENCE_MODEL_NAME)
self.set_prompt_template(PROMPT_TEMPLATE)
# modelのパラメータ
self.max_new_tokens = 64
self.temperature = 0.7
self.repetition_penalty = 1.1
self.top_p = 0.9
self.repetition_penalty = 1.05
def load_model(self, model_name: str):
print(f">> load inferencing model from huggingface: {model_name}")
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cpu")
if torch.cuda.is_available():
self.model = self.model.to("cuda")
def inference(self, question: str, context: list):
input_prompt = self.generate_prompt(question, context)
inputs = self.tokenizer(input_prompt, return_tensors="pt").to(self.model.device)
with torch.no_grad():
tokens = self.model.generate(
**inputs,
max_new_tokens=self.max_new_tokens,
do_sample=True,
temperature=self.temperature,
top_p=self.top_p,
repetition_penalty=self.repetition_penalty,
pad_token_id=self.tokenizer.pad_token_id,
)
output = self.tokenizer.decode(tokens[0], skip_special_tokens=True)
return output
一連の処理をまとめる
VectorDBクラスとInferencingModelクラスに実装した2.~5.の処理を取りまとめるInferenceLogicクラスを作成します。クラス内のinference()は引数のinput_questionにユーザが直接入力した質問(prompt)受け取り、input_questionを元にDBのベクトル検索、その結果を用いて推論を行うものになります。
from database.vector_db import VectorDB
from model.inferencing_model import InferencingModel
class InferenceLogic:
def __init__(self):
self.db = VectorDB()
self.model = InferencingModel()
def inference(self, input_question: str):
context_list = self.db.search_similary_sentences(input_question)
output = self.model.inference(input_question, context_list)
return output
if __name__ == "__main__":
rag_app = InferenceLogic()
question = "RAGについておしえて！"
output = rag_app.inference(question)
print(output)
おわりに
RAGをLangChainを使わずに作ってみる（理解編）と（実装編）は以上になります。
RAGの一連の処理を自ら実装することで、RAGを用いた推論精度を高めるため要所が明らかになりとてもよかったと思います。
また、今回の記事はGPU環境がなくても動作するようなモデルを選定しているので、手元で動かしてもらえると嬉しいです。