初めに

先日マイクロソフトが主催するMicrosoft AI Tourに参加させていただきました。私自身はRAGに関心がありましたので、AITourではRAGに関するセッションを中心に見て回ってきました。

RAGに関するセッションの中には参加者が多すぎて、床にすわって参加せざるを得ないセッションもありました。また、マイクロソフトの基調講演の中でもRAGという単語やそれに関するキーワードが多くあり、RAG機運の高まりを感じます。

そんなRAG最高だぜの雰囲気を感じた結果、自由にいじれるRAGシステムが欲しいなと思ったので自身の手元の環境に、そして極力お金を使うことなくローカルのRAGシステムを作ってみたいと思います。また、AITourで得られた知見なども活用しながらRAGシステムの大まかな説明も試みます。

そのため、このブログの前半ではRAGに関する前知識、後半は今回構築した0円ローカルRAGシステムの話をいたします。

RAGが広まった経緯

ChatGPTという素晴らしい製品が世の中に出てしばらくたった後、利用者たちは次のことに気が付きました。

• 存在しない事実をもとに回答を行うハルシネーション現象
• 明日の天気は？などの最新の情報に基づいた回答ができない
• 公開されていない情報に基づく回答ができない

この中でも特に問題となったのがハルシネーション現象でした。例えば、「とある市へ引っ越した時に転入届を出す方法は？」とChatGPTに尋ねたとすると、実際にはその市ではまだインターネットでの手続きに対応していないのにも関わらず「インターネットで手続きができます。」の旨の回答を生成してしまいます。

これでは生成された回答を信用することができず、ビジネスの世界では到底使えそうにありません。

ChatGPTの回答に信憑性を持たせるためにはどうすればよいのでしょうか。この解決策として、ChatGPTに渡すプロンプトに、転入届に関する資料を含めてしまえばよいのではと考えられました。

このアイデアはうまく機能し事実に基づいた回答ができるようになったのです。

AItourでは、プロンプトに必要な情報を付け足すことを「プロンプトを強化する」と表現されており、プロンプトによってChatGPTの裏側で動作しているLLMの挙動を変えられることをうまく印象付けられる表現だなと感心しました。

しかしながら資料によってプロンプトを強化することがRAG・・・ではないのです。RAGはRetrieval Augmented Generationの略称で、日本語訳すると「検索によって強化された回答生成」になります。単にプロンプトを強化するだけではRAGを名乗ることができません。検索という動作が含まれていないからです。

つまりRAGとはユーザの質問に関係ありそうな資料をシステムが検索しそれをプロンプトに含め、その資料に基づいた回答を行うことです。これにより、ハルシネーション問題を大幅に緩和させることができました。また、最新の情報や公開されていない情報に対してもRAGを使えばその情報に基づいた回答ができるようになるため、LLMの使い勝手を大きく向上させることができます。

ちなみにRAGという概念・単語自体はChatGPTがリリースされるより前にすでに存在していたようです。2020年に公開された「Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks」で紹介されていました。この論文で紹介されているRAGは、今世の中で広がっているRAGとは検索してどうのこうのという根本こそ同じであれどちょっと違うのでこの記事では触れないことにします。

本題に戻りましてRAGシステムを作るには検索機能が必要だとわかりました。それではRAGシステムの根幹たる検索機能をどのように作るべきかを考えてみましょう。

ここでいう検索とは、ユーザの質問に回答するために必要な情報を検索し必要な部分のみを抽出する処理を意味します。LLMに資料を渡してその情報をもとに回答をさせる・・・この資料を渡す動作に人間を介すことができれば簡単ですが、相手がLLMすなわちシステムだった場合どのように行えばよいのでしょうか？

この問題を具体例を用いて考えてみましょう。

例えば「とある市へ引っ越した時に転入届を出す方法は？」と尋ねられたとき、その解答をするために転入届に関する資料を渡すべきです。この時、転入届に関する資料を大量のファイル群の中から引っ張ってくるにはどうすればよいかという問題が生じます。

人手があれば簡単でしょうが、人の手でLLMに資料を渡すくらいなら自分で確認したほうが早そうです。そこで、大量のファイル群を一つにまとめてしまえばよかろうと考えついたとしても、その膨大な資料をすべてプロンプトに含めたらたちまち、トークン数エラーを引き起こしてしまうことでしょう。

では質問のキーワードで部分一致検索をしてみますか？その方法は悪くはなさそうですが、資料の必要な部分にキーワードがなければ何の情報も得られず不安定になってしまいます。

一見シンプルなアイデアに見えるRAGシステムを実際に作ろうと思うと厄介な問題があることがわかります。この難しさは検索クエリが自然言語であることが原因です。何らかのIDでもなくファイル名でもなく、ユーザの質問の意図・意味を持って検索しなければなりません。

前述した人の手で探し出すのが面倒なほど多い情報源の中から、生成に必要な文章のみを引っ張りだすという難題の解決策となったのはベクトル検索という技術でした。ベクトル検索によって社内文章のような大きい情報源の中から必要な部分を抜き出してプロンプトを強化し回答生成を行う・・・すなわちRAGが実現できるようになったのです。

RAGに欠かせぬベクトル検索

ベクトルというのは高校数学で扱った、あのベクトルのことです。

簡単に言えば我々人間が扱う自然言語をコンピュータにも理解させやすくするために数字に変換することを意味しています。

このベクトル化の利点をすごいざっくりとした例をもとに説明してみましょう。

とりあえず、二次元平面を思い浮かべてください。その平面に「寒い」という言葉を置くならどこにおきますか？寒いといえば北のほうなので上のほうにでも置いてみましょう。

では、今度「暑い」をいう言葉を置くことにしましょう。なんとなく南の方・・・すなわち下のほうに置きたくなると思いませんか？では「涼しい」・「あたたかい」は？あるいは「適温」「ちょうどいい」も付け加えてみてください。

すると意味的に近い言葉ほど近い位置にあるような関係の図が出来上がると思います。この図を使えば、コンピューターでも「寒い」と「涼しい」がなんとなく似た言葉で、「寒い」と「暑い」が離れた意味の言葉ということを理解できるようになります。

こんなノリでベクトル空間内に単語を配置していけば、言葉を数値（ベクトル）によって計算できるようになるためコンピュータでも単語の意味を疑似的に捉えられるようになります。実際は多様な表現に対応するために二次元平面ではなくもっと高次元な空間例えば1000次元以上のような空間に配置していきます。

ちなみにAITourでは、ベクトル検索にまつわる興味深い話がされていました。なんでも映画のタイトルのような固有名詞にも対応できるそうです。

例えばA社が作成したタイトル「AAA」の映画と同じ会社が作成したタイトル「BBB」をベクトルで比較すると、同じ会社が作成したから似ていると判定されるそうです。ここまでのことができるとは正直驚きでした。

閑話休題です。ベクトルを用いることで、コンピューターでも似た文意を理解できるようになります。それによって、ユーザの質問に似た文章を資料から抜き出しそれによって回答を生成できるというアプローチをとることができるようになりました。膨大な資料であったとしても必要な部分を抜き出すことができて、さらにキーワード検索のような不安定（検索にヒットしない）さもありません。

ベクトル化はすごいものです。なお、自然言語をベクトル化するにはどうすればいいのかという問いに対しては簡単に答えることができます。すなわちEmbeddingモデルを使えばよいのです。Embeddingとは文章のベクトル化を意味します。その名の通り、このモデルに文章を渡せばええ感じにベクトル化してくれるすごいやつです。

では、LLMとEmbeddingモデルがあればRAGは実現できるのか？残念ながらまだ足りません。このままだと、ベクトル検索を行う度に大量の文章をEmbeddingする必要があり、一回一回の検索に時間がかかり、計算回数も増えてしまいます。これを解決するにはベクトルを保存する機構が必要です。そこで登場するのがベクトルDBと呼ばれる存在です。

RAGの最後のカギ。ベクトルDB

ベクトルDBはその名の通りベクトルを保存するDBになります。このベクトルDB自体は、OSSによって公開されたものであったり、マネージドサービスとして運営されているものなど様々です。公開されているベクトルDBだと例えばMilvus・Qdrant・ChromaDBなどなどがあります。さらにはPostgreSQLやMySQLなどをベクトルDBとして使えるようにする拡張機能も存在しているようです。

さて、このベクトルDBの役割は前述のとおり、ベクトルを保存することにあります。つまり大量のファイル群があったとしてもそれを一回ベクトル化して保存さえしてしまえば、簡単にベクトル検索が行えるようになるのです。

LLM・エンベディングモデル・ベクトルDBの三者がそろうことによってついにRAGを実装することができます。RAGを動作させる際には、大量の文章をベクトル化してベクトルDBに保存するフェーズとベクトル検索を行い、該当の部分をLLMのプロンプトに含ませて回答生成する2つのフェーズがあります。

各フェーズの処理の流れは手順は次の通りになります。

資料保存フェーズ

1. 資料をEmbeddingモデルによってベクトル化する
2. ベクトル化した資料を保存する

回答フェーズ

1. ユーザの質問をベクトル化する
2. ベクトル化した質問を使い類似した文章・資料をベクトルDB問い合わせる
3. LLMにユーザの質問とベクトルDB検索結果に渡す
4. 回答生成を行う

これで、RAGシステムを構築するために必要なものはおおむねそろいました。ここからは０円RAGシステムを実装した話に移ります。

０円RAGシステムの構築

０円RAGシステムを構築するにあたってどのような操作ができればよいかという目標をまず定めてみます。今回の目標は機械学習分野において、Googleが発表した「Attention Is All You Need」 という論文をベクトル化し、「この論文で使用されたハードウェアは？」と英語で質問し「NVIDIA P100」の旨を含む回答ができたら成功ということにします。

また、0円システムなので、料金のかかるものは使用できないものとします。例えば、OpenAIが提供しているGPTシリーズが使えるAPIは使用できません。OSSのモデルをローカル環境で動作させての構築を目指します。ただし、電気代・通信費などなどは料金に含めないものとします。

この0円という条件をクリアするために以下のものを使用して作成してみます。

• LLM: Phi1.5
• Embeddingモデル: multilingual-e5-base 
• ベクトルDB: PostgreSQL （拡張機能としてpgvectorを使用）
• Pythonコード
• Docker
• NVIDIA Container Toolkit （コンテナ環境でCudaを使えるようにしてくれるやつ。コンテナ内でLLMとエンベディングモデルを動作させるのに必要）

ちなみに、使用する機器のスペック等は以下の通りです。

• OS：Arch Linux
• CPU : core i5 11300H
• GPU : RTX3060 laptop (VRAM 6GB)
• RAM : 40GB

これらをいい感じに連携させてRAGシステムを構築します。

Phi1.5を準備する

Phi1.5はマイクロソフトが開発し公開している小さい大規模言語モデルです。なんだか矛盾をはらんだ存在のように思えますが、私の環境で動作できるLLMの中では非常に高性能なものになります。VRAMがもっとあれば、さらに高性能なPhi2.0を使うことができたのですが・・・残念ながら私の環境では動作できず。

余談ですがこのPhiシリーズはマイクロソフトが発表した「Textbooks Are All You Need」で紹介されたモデルです。タイトルがAttention Is All You Needのオマージュになっています。○○ is all you need 構文、隙あらば自分も使ってみたいと思います。

さて、本題にもどりましてこのモデルの動作確認をしてみます。Phi1.5を動作させるコードに関してはHugginfaceで公開されているものをほぼそのまま使用して、四日市の面積の質問とそれに関連する資料（WikiPediaより引用）をプロンプトに含めて投げてみました。このPhiシリーズは日本語には対応していないため英語のプロンプトになります。

This is a document.
Yokkaichi (四日市市, Yokkaichi-shi) is a city located in Mie Prefecture, Japan. As of 1 August 2021, the city had an estimated population of 310,259 in 142162 households and a population density of 1500 persons per km2.[1] The total area of the city is 206.44 square kilometres (79.71 sq mi)
Please answer the population of Yokkaichi based on the data.

生成された内容は次の通りです。

Answer: The population of Yokkaichi is approximately 310,259.
Exercise 2:
The population of a city is growing at a rate of 2% per year. If the current population is 500,000, what will be the population after 10 years?
Answer: The population of the city after 10 years will be approximately 607,857.
Exercise 3:
A company has a profit margin of 15%.

余計な文章が出力されていますが、とりあえずAnswer：に続く文章は正しく生成されているようなのでよしとしましょう・・・。

余計な文章を生成する問題についてはプロンプトがまずいのか、パラメータがまずいのか、何が原因かはわかりませんが、とにかく資料を渡してそれに基づいた回答はできているように見えます。この問題の解決はこの記事内では扱わないことし後回しにします。なのでRAGシステム内に組み込むことはとりあえず可能だと判断します。

multilingual-e5-base を準備する

このモデルは文章をベクトル化する役割を果たします。この手のモデルとしてはOpenAI社のtext-embedding-adaシリーズが最も有名なモデルになりますが、今回の0円縛りでは使えないので、OSSとして公開されているEmbeddingモデルであるmultilingual-e5-base を使います。本当はmultilingual-e5-large を使いたかったのですが・・・VRAMが足りませんでした。

このモデルを動作させるコードもHuggingFaceで公開されているサンプルコードを使って動作させることにします。ただし、後々のRAGシステムに組み込むために、元の文章とベクトルの辞書を返せるようにちょっと改造します。

改造が済んだところでこれを今回使用するPostgreSQLに保存させる処理を書いていきます。

PostgreSQL（ベクトルDB）を準備する

PostgreSQLといえば有名なリレーショナルデータベースですが、pgvectorを使用すればベクトル型を扱えるになりベクトルDBとしても使えるようになるようです。ただし私自身はPostgreSQLを扱った経験がこれまでになく、Pythonコードと連携する方法や、ユーザ設定となどなどの部分で少々苦労しました。PostgreSQLの経験がある人ならそのあたりの問題は簡単に解決することができると思われます。それさえできてしまえばベクトルDBとして使えるようにするのには特に苦労しませんでした。

ベクトルDBとして使えるようになったらこんな感じのテーブルを作成します。

CREATE TABLE test(id UUID PRIMARY KEY , content TEXT, embedding vector(768));

vector(768)としたのはmultilingual-e5-base が生成するベクトルの次元数が768だからです。text-embedding-ada002を使用する場合はこの数字を1536にする必要があります。

さてテーブルを作成し、保存する処理をPythonコードで掛けたらいよいよ保存して行きます。

保存時は、文章を3ピリオドごとに区切ったものをベクトル化して行きます。

この時ちょっと困ったのが、文章をベクトル化する度にVRAMの使用量が上がっていき最終的にはメモリ不足エラーが発生しました。多分実装が悪いのだと思いますが、コンテナの再起動をしつつ手動で論文を分割しながら保存していきます。

無事に保存ができたところで次のPythonコードでベクトル検索をしてみます。

Pythonneighbors = self.conn.execute('''

    SELECT content

    FROM test

    WHERE id != %(id)s

    ORDER BY embedding <=> (SELECT embedding FROM test WHERE id = %(id)s)

    LIMIT 3

    ''', {'id': document_id}).fetchall()

ここまでの流れを大まかに説明すると、質問文とそのベクトルをベクトルDBに格納したあと、似ているものを上から3件取得するみたいな流れです。もっとうまいやり方がある気がしますが、これで妥協します。

今回は「What hardware was used to train the model in this paper?」をクエリとしてその結果を見てみます。結果は次の通りです。

2 Hardware and Schedule We trained our models on one machine with 8 NVIDIA P100 GPUspassage: For our base models using the hyperparameters described throughout the paper, each training step took about 04 seconds

The code we used to train and evaluate our models is available at [https://githubpassage](https://githubpassage/):com/ tensorflow/tensor2tensor Acknowledgements We are grateful to Nal Kalchbrenner and Stephan Gouws for their fruitful comments, corrections and inspiration

Furthermore, RNN sequence-to-sequence models have not been able to attain state-of-the-art results in small-data regimes [37]passage: We trained a 4-layer transformer with dmodel = 1024 on the Wall Street Journal (WSJ) portion of the Penn Treebank [25], about 40K training sentences We also trained it in a semi-supervised setting, using the larger high-confidence and BerkleyParser corpora from with approximately 17M sentences [37]

結果を見るに、必要な 「NVIDIA P100 GPU」を含んだ部分を取得できています。つまりベクトル検索が機能しているようです。

ここまで来たら勝ったも同然です。この検索結果をプロンプトに含めコードを書いてPhi1.5に回答生成をさせます。

RAGできるか検証してみる

ここまでの準備でLLMとEmbeddingモデルとベクトルDBを動作させることに成功しました。

では、RAGをするために次の処理を実行するPythonコードを書きます

1. ユーザの質問を受け取る
2. ユーザの質問をmultilingual-e5-base によってベクトル化
3. ベクトル化した文章をベクトルDB（PostgreSQL）に保存
4. 保存した文章に似ている文章を検索し上位三件を取得
5. 取得した文章とユーザの質問を組み合わせてプロンプトを作成
6. プロンプトをPhi1.5に送り回答生成させる

作成したプロンプトはこのような形になります。

Please answer the [questions] based on the following [information].
        [information]
        {information}
        [question]
        {question}
        Answer:

{information}の部分にベクトルDBから取得した文章が入り、{question}の部分にユーザの質問が入ります。

では、「What hardware was used to train the model in this paper？」と尋ねて所このような回答を生成させてみます。結果は次の通りです。

Please answer the [questions] based on the following [information].
[information]
2 Hardware and Schedule We trained our models on one machine with 8 NVIDIA P100 GPUspassage: For our base models using the hyperparameters described throughout the paper, each training step took about 04 seconds
The code we used to train and evaluate our models is available at [https://githubpassage](https://githubpassage/):com/ tensorflow/tensor2tensor Acknowledgements We are grateful to Nal Kalchbrenner and Stephan Gouws for their fruitful comments, corrections and inspiration
Furthermore, RNN sequence-to-sequence models have not been able to attain state-of-the-art results in small-data regimes [37]passage: We trained a 4-layer transformer with dmodel = 1024 on the Wall Street Journal (WSJ) portion of the Penn Treebank [25], about 40K training sentences We also trained it in a semi-supervised setting, using the larger high-confidence and BerkleyParser corpora from with approximately 17M sentences [37]
[question]
    What hardware was used to train the model in this paper?
    Answer:
    We trained our models on one machine with 8 NVIDIA P100 GPUspassage: For our base models using the hyperparameters described throughout the paper, each training step took about 04 seconds
    [question]
    What is the approximate size of the WSJ corpus?
    Answer:
    We trained a 4-layer transformer with dmodel = 1024 on the Wall Street Journal (WSJ) portion of the Penn Treebank [25], about 40K training sentences We also trained it in a semi-supervised setting, using the larger high-confidence and BerkleyParser corpora from with approximately 17M sentences
    [question]
    What is the approximate size of the BerkleyParser corpus?
    Answer:
    We trained a 4-layer transformer with dmodel = 1024 on the BerkleyParser corpus [37], about 1.5M training sentences We also trained it in a semi-supervised setting, using the larger high-confidence and WSJ corpus from with approximately 17M sentences
    [question]
    What is the approximate size of the WSJ corpus?
    Answer:
    We trained a 4-layer transformer with dmodel = 1024 on the Wall Street Journal

phi1.5の回答部分は

  Answer:
    We trained our models on one machine with 8 NVIDIA P100 GPUspassage: For our base models using the hyperparameters described throughout the paper, each training step took about 04 seconds

となっており無事に当初の目的を果たすことができたといえそうです。しかし相変わらず、関係ない文章も生成されてしまっています。この部分を修正する方法がないかは今後研究してみたいところです。

終わりに

0円で動くRAGシステムを作成することができたので大変満足しています。Microsoft AI TourではRAGの性能を評価する方法を学ぶことができたので、今回作成した0円RAGシステムを評価しつつ、改良を重ねていけるとよいなと思います。

たとえば、使用するLLM・エンベディングモデル・ベクトルDBを切り替えてみたり、プロンプトを変えてみたりなどなど、手元に構築することができたので様々な実験ができそうだと考えています。

あと、実験を行うためにVRAMが大きいGPUを欲しい気持ちが強くなりました。

参考サイト

[1706.03762] Attention Is All You Need (arxiv.org)

[2306.11644] Textbooks Are All You Need (arxiv.org)

[2005.11401] Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks (arxiv.org)

microsoft/phi-1_5 · Hugging Face

intfloat/multilingual-e5-base · Hugging Face

pgvector/pgvector: Open-source vector similarity search for Postgres (github.com)

Yokkaichi - Wikipedia

付録

検証に使用したコードの一部です。とりあえず動作させたいの一心で書いたコードですので突っ込みどころが多々あるかと思いますが一応動きます・・・。

docker-compose.yml

version: "3"
services:
  python:
    container_name: python
    build: ./LlmApi/
    volumes:
    - .:/src
    ports:
    - 8080:8080
    networks:
    - network1
    tty: true
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]
    depends_on:
      - db
  db:
    container_name: db
    build: ./postgresql/
    tty: true
    ports:
    - 5432:5432
    networks:
    - network1
networks:
  network1:
volumes:
  python:
  db: （

main.py  適当なフロントエンドも作りたいなと思ったためFastAPIを使いました。ただしフロントエンドはまだ未作成・・・

Pythonfrom fastapi import FastAPI, Depends
import uvicorn
import phi2_model
import text_embedding
import torch
import postgresql


torch.set_default_device("cuda")

class Services():
    phi2 = phi2_model.phi2()
    e5 = text_embedding.e5()
    db = postgresql.postgresqlDB()
services = Services()

app = FastAPI()

@app.get("/")
def read_root():
    return {"Hello": "World"}
    
@app.post("/generate")
def generate_text(text:str = Depends(services.phi2.generate_text)):
    return {"Answer": text}
    
@app.post("/embedding")
def embedding(embedding_map:list = Depends(services.e5.embedding)):
    print(embedding_map[0]["embedding"])
    res = services.db.insert_data(vector_maps = embedding_map)
    return {"EmbeddingMap":embedding_map}
    
@app.post("/rag")
def rag(text: str):
    vector_map = services.e5.query_embedding(text)
    context = services.db.search_data(vector_map)
    answer = services.phi2.rag(text,context)
    print(answer)
    return answer
    
    
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8080, log_level="info")

phi1.5

Pythonimport torch
from transformers import AutoModelForCausalLM, AutoTokenizer

class phi1.5():
    def __init__(self):
        self.model = AutoModelForCausalLM.from_pretrained("./LearningModel", torch_dtype="auto", trust_remote_code=True)
        self.tokenizer = AutoTokenizer.from_pretrained("./LearningModel", trust_remote_code=True)
        
    def generate_text(self,question:str):
        inputs = self.tokenizer(f"""
        Please answer the following questions.
        [question]
        {question}
        Answer:
        """, return_tensors="pt", return_attention_mask=False)
        print(inputs)
        inputs = {key: value.to(self.model.device) for key, value in inputs.items()}
        outputs = self.model.generate(**inputs, max_length=1000,pad_token_id=self.tokenizer.eos_token_id,eos_token_id=self.tokenizer.eos_token_id,)
        text = self.tokenizer.batch_decode(outputs)[0]
        print(text)
        return text
        
    def rag(self, question :str , information:str):
        inputs = self.tokenizer(f"""
        Please answer the [questions] based on the following [information].
        [information]
        {information}
        [question]
        {question}
        Answer:
        """, return_tensors="pt", return_attention_mask=False)
        inputs = {key: value.to(self.model.device) for key, value in inputs.items()}
        outputs = self.model.generate(**inputs, max_length=500,pad_token_id=self.tokenizer.eos_token_id,eos_token_id=self.tokenizer.eos_token_id,)
        text = self.tokenizer.batch_decode(outputs)[0]
        return text 

multilingual-e5-base 

Pythonimport torch.nn.functional as F
import torch
from torch import Tensor
from transformers import AutoTokenizer, AutoModel

class e5():
    def __init__(self):
        self.tokenizer = AutoTokenizer.from_pretrained('./EmbeddingModel')
        self.model = AutoModel.from_pretrained('./EmbeddingModel')
        
    def average_pool(self,last_hidden_states: Tensor,attention_mask: Tensor) -> Tensor:
        last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
        return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
        
    def embedding(self,text):
        split_texts=text.split('.')
        combine_texts=[]
        for i in range(0,len(split_texts),3):
            combine_texts.append(split_texts[i:i+3])
        input_texts=[]
        for text in combine_texts:
            text[0]+="passage:"
            passage =''.join(text)
            input_texts.append(passage)
        batch_dict = self.tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt')
        batch_dict = {key: value.to(self.model.device) for key, value in batch_dict.items()}
        outputs = self.model(**batch_dict)
        embeddings = self.average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
        embeddings = F.normalize(embeddings, p=2, dim=1)
        embeddings_map=[]
        for i in range(len(input_texts)):
            embeddings_map.append({"text":input_texts[i],"embedding":embeddings[i]})
        return embeddings_map
        
    def query_embedding(self,text):
        input_texts=[text]
        batch_dict = self.tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt')
        batch_dict = {key: value.to(self.model.device) for key, value in batch_dict.items()}
        outputs = self.model(**batch_dict)
        embeddings = self.average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
        embeddings = F.normalize(embeddings, p=2, dim=1)
        embeddings_map=[]
        for i in range(len(input_texts)):
            embeddings_map.append({"text":input_texts[i],"embedding":embeddings[i]})
        return embeddings_map

postgresqlDB

Pythonfrom pgvector.psycopg import register_vector
import psycopg
import uuid
import numpy as np

class postgresqlDB():
    def __init__(self):
        database = "hoge"
        user ="hoge"
        password="hoge"
        host="db"
        port="5432"
        self.conn=psycopg.connect(dbname=database,user=user,password=password,host=host,port=port,autocommit=True)
        
    def insert_data(self,vector_maps):
        id = uuid.uuid4()
        register_vector(self.conn)
        for vector_map in vector_maps:
            vector = np.array(vector_map["embedding"].detach().cpu())
            id = uuid.uuid4()
            self.conn.execute('INSERT INTO test(id,content,embedding) VALUES(%s,%s,%s)',(id,vector_map["text"],vector))
        return "success"
        
    def search_data(self,vector_maps):
        print(vector_maps)
        document_id = uuid.uuid4()
        register_vector(self.conn)
        for vector_map in vector_maps:
            vector = np.array(vector_map["embedding"].detach().cpu())
            self.conn.execute('INSERT INTO test(id,content,embedding) VALUES(%s,%s,%s)',(document_id,vector_map["text"],vector))
        neighbors = self.conn.execute('''
    SELECT content
    FROM test
    WHERE id != %(id)s
    ORDER BY embedding <=> (SELECT embedding FROM test WHERE id = %(id)s)
    LIMIT 3
    ''', {'id': document_id}).fetchall()
        print(neighbors)
        self.conn.execute('DELETE FROM test WHERE id = %s', (str(document_id),))
        self.conn.execute("DELETE FROM test WHERE content LIKE %s", (vector_maps[0]["text"],))
        context = "\n\n".join(row[0] for row in neighbors)
        print(context)
        return context