RAGをLangChainを使わずに作ってみる（実装編）

# はじめに
今回は最近、流行りのRAG(Retrieval-Augmented Generation)をLangChainなどを使用せずに処理を構築した内容を紹介する実装編になります。前回は理解編としてRAGの処理の流れを紹介しました。良ければ前回の記事と合わせてみて頂けると幸いです。

前回の記事: [RAGをLangChainを使わずに作ってみる（理解編）](https://qiita.com/DeOma004/items/0ab1c8afa0667ee71ac2)

# RAGを実装する
## 処理の概要

RAG の処理は以下の 2 つの処理に分けることができ、それぞれの大まかな処理は以下でした。

- Inferencing: RAGをつかったLLMの推論
    1. LLMモデルに質問(prompt)の投げる
    2. 質問(prompt)を埋め込みモデルでベクトルに変換する
    3. 変換したベクトルをもとにDB内を検索する
    4. 検索結果の引用情報を質問(prompt)に埋め込む
    5. 引用情報を埋め込んだ質問(prompt)を用いて推論を行う

- Ingestion: ベクトルDBの準備
    1. 関連文書を任意のサイズ(chunk_size)で分割する
    2. 関連文書を埋め込みベクトルに変換する
    3. DBに関連文書とベクトルを記録する

処理を図示すると以下のように表現でき、赤く囲った部分がIngestion処理、青く囲った部分がInferencing処理となります。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/682846/2732c525-6fcf-b36a-882c-15ec6b723ec0.png)

今回はこれらの処理を実装する主なライブラリしてモデルまわりはhugging face関連のものを利用し。DBはelastic searchを用い、dockerを用いて環境の構築を行います。

# 使用する環境について
## ベクトルDBに関して

### elastic searchコンテナの立ち上げ
ベクトル DB は elastic search を使用し関連文書の検索、ベクトルの管理を行います。
DBは以下のdockerを用いて環境を構築します。docker-composeではデータの確認用としてkibanaも合わせて起動させています。
elasticsearchが起動後に`http://localhost:5601/app/home#/`へアクセスしてデータの確認ができると思います。

__dockerfile__
```dockerfile
FROM docker.elastic.co/elasticsearch/elasticsearch:8.1.0
RUN elasticsearch-plugin install analysis-kuromoji
```

__docker-compose.yml__
```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（文書）の保存情報の定義は今回は以下のようにしました。

```python
mapping_info = {
    "mappings": {
        "properties": {
            "file_title":{
                "type": "text"
            },
            "context":{
                "type": "text"
            },
            "context_vector":{
                "type": "dense_vector",
                "dims": 384
            }
        }
    }
}
```

Pythonでindexを登録する処理は以下になります。
```python
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](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2)
- 推論モデル：[cyberagent/open-calm-small](https://huggingface.co/cyberagent/open-calm-small)

# Ingestion
まずは既存文書を埋め込みベクトルへ変換し、DBに保存する処理について紹介します。Ingestionで必要な処理は以下でした。

- Ingestion: ベクトルDBの準備
    1. 関連文書を任意のサイズ(chunk_size)で分割する
    2. 関連文書を埋め込みベクトルに変換する
    3. DBに関連文書とベクトルを記録する

ここでは`VectorDB`クラスを定義し、処理を実装していきます。
なお、細かいDBでのデータ管理は考慮していないので、文書のベクトルを扱う最低限の処理の実装の紹介になります。

## 1. 関連文書を任意のサイズ(chunk_size)で分割する
ここではベクトルに変換する既存文書を任意のサイズ(`chunk_size`)で分割する処理になります。

`chunk_size`で分割する意図としては検索する情報量を調整するためになります。検索結果として文書全体が返ってきたとしても、具体的にどの部分を引用すれば良いかがわからないと思います。なので、検索結果として適切な粒度となるように分割し、埋め込みベクトルに変化します。

`chunk_size`としては文書のページ単位や段落単位など様々な粒度が考えられます。今回は文字数で情報を区切ることとしました。

```python
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`を使用して読み込んだモデルを返す関数を実装しています。

必要なコードとしては以下の内容になります。

```python
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`メソッドを呼ぶだけになります。

```python
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が存在しない場合に、作成する処理を組み込んでいます。


```python
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に記録されます。

```python
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`を用いて複数のデータを記録する処理になっています。記録方法の詳細に関しては[ココ](https://qiita.com/satto_sann/items/8a63761bbfd6542bb9a2#%E3%83%90%E3%83%AB%E3%82%AF%E3%82%A4%E3%83%B3%E3%82%B5%E3%83%BC%E3%83%88)を参考にしました。


```python
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に情報を登録できます。

```python
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の推論
    1. LLMモデルに質問(prompt)の投げる
    2. 質問(prompt)を埋め込みモデルでベクトルに変換する
    3. 変換したベクトルをもとにDB内を検索する
    4. 検索結果の引用情報を質問(prompt)に埋め込む
    5. 引用情報を埋め込んだ質問(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へのクエリ
    ```python
    {
        "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しています。

```python

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の部分がユーザーからの入力となるように処理しています。

```python
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のモデルカードに記載されているので、その部分を参考に実装してみてください。

```python
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のベクトル検索、その結果を用いて推論を行うものになります。

```python
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環境がなくても動作するようなモデルを選定しているので、手元で動かしてもらえると嬉しいです。