pgvectorを試す

docker run --rm -d \
    --name pgvector-test \
    -p 5432:5432 \
    -e POSTGRES_PASSWORD=postgres \
    pgvector/pgvector:pg17

docker exec -ti pgvector-test bash

psql -U postgres 

psql (17.0 (Debian 17.0-1.pgdg120+1))
Type "help" for help.

postgres=#

CREATE DATABASE vectordb;

\c vectordb

CREATE EXTENSION vector;

CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3));

INSERT INTO items (embedding) VALUES ('[1,2,3]'), ('[4,5,6]');

SELECT * FROM items ORDER BY embedding <-> '[3,1,2]' LIMIT 5;

 id | embedding
----+-----------
  1 | [1,2,3]
  2 | [4,5,6]
(2 rows)

mkdir pgvector-test && cd pgvector-test

docker run --rm \
    -p 8888:8888 \
    -u root \
    -e GRANT_SUDO=yes \
    -v .:/home/jovyan/work \
    quay.io/jupyter/minimal-notebook:latest

!pip install pandas psycopg2-binary openai

import psycopg2

conn = psycopg2.connect(
    user='postgres',
    password='postgres',
    database='vectordb',
    host='host.docker.internal',
    port=5432     
)

cur = conn.cursor()

cur.execute("SELECT * FROM items ORDER BY embedding <-> '[3,1,2]' LIMIT 5")

records = cur.fetchall()

print(records)

cur.close()
conn.close()

[(1, '[1,2,3]'), (2, '[4,5,6]')]

!wget https://tulip.kuee.kyoto-u.ac.jp/localgovfaq/localgovfaq.zip
!unzip localgovfaq.zip

!pip install pandas

import pandas as pd

def file2list(filename: str, prefix: str = "") -> tuple:
    """Q/AファイルをIDとコンテンツに分割、それぞれを配列で返す"""
    contents = []
    ids = []
    try:
        with open(filename, 'r') as file:
            for line in file:
                line = line.strip().replace(" ", "")
                id, content = line.split('\t')
                if prefix:
                    id = f"{prefix}_{id}"
                contents.append(content)
                ids.append(id)
    except Exception:
        raise

    return contents, ids

questions, ids = file2list("localgovfaq/qas/questions_in_Amagasaki.txt")
answers, _ = file2list("localgovfaq/qas/answers_in_Amagasaki.txt")

raw_df = pd.DataFrame({'QA_ID': ids, 'QUESTION': questions, 'ANSWER': answers})
df = raw_df.head(100)
df

import getpass
import os

os.environ["OPENAI_API_KEY"] = getpass.getpass('OPENAI_API_KEY')

from openai import OpenAI
from tqdm import tqdm

tqdm.pandas()

client = OpenAI()

def get_embedding(text):
    response = client.embeddings.create(
        input=text,
        model="text-embedding-3-small"
    )
    return response.data[0].embedding

df['embedding'] = df['ANSWER'].progress_apply(get_embedding)
df

import psycopg2

create_table_query = """
    CREATE TABLE IF NOT EXISTS faq (
        id bigserial PRIMARY KEY,
        qa_id int,
        answer text,
        embedding vector(1536)
    );
"""

with psycopg2.connect(
    user='postgres',
    password='postgres',
    database='vectordb',
    host='host.docker.internal',
    port=5432
) as conn:
    with conn.cursor() as cur:
        cur.execute(create_table_query)
        conn.commit()

docker exec -ti pgvector-test bash

psql -U postgres vectordb

 \dt

         List of relations
 Schema | Name  | Type  |  Owner
--------+-------+-------+----------
 public | faq   | table | postgres
 public | items | table | postgres
(2 rows)

 \d faq

                                Table "public.faq"
  Column   |     Type     | Collation | Nullable |             Default
-----------+--------------+-----------+----------+---------------------------------
 id        | bigint       |           | not null | nextval('faq_id_seq'::regclass)
 qa_id     | integer      |           |          |
 answer    | text         |           |          |
 embedding | vector(1536) |           |          |
Indexes:
    "faq_pkey" PRIMARY KEY, btree (id)

import psycopg2
from psycopg2.extras import execute_values

with psycopg2.connect(
    user='postgres',
    password='postgres',
    database='vectordb',
    host='host.docker.internal',
    port=5432
) as conn:
    with conn.cursor() as cur:
        for _, row in tqdm(df.iterrows(), total=len(df)):
            insert_query = "INSERT INTO faq (qa_id, answer, embedding) VALUES (%s, %s, %s)"
            cur.execute(insert_query, (row['QA_ID'], row['ANSWER'], row['embedding']))
        conn.commit()

query_text = "地域総合センター今北へはどう行けばいいですか？"
query_embedding = get_embedding(query_text)

# PostgreSQLでの類似検索
with psycopg2.connect(
    user='postgres',
    password='postgres',
    database='vectordb',
    host='host.docker.internal',
    port=5432
) as conn:
    with conn.cursor() as cur:
        search_query = """
            SELECT qa_id, answer, embedding <=> %s::vector AS distance
            FROM faq
            ORDER BY distance
            LIMIT 5;
        """
        cur.execute(search_query, (query_embedding, ))
        results = cur.fetchall()

for result in results:
    qa_id, answer, distance = result
    print(f"Distance: {distance}, QA_ID: {qa_id}, Answer: {answer}")
    print("-----")

Distance: 0.40620364528286224, QA_ID: 1, Answer: ■地域総合センター今北には、十分な駐車場がございませんので、市バスをご利用ください。ＪＲ沿線からは「立花駅」、阪急沿線からは「塚口駅」「武庫之荘駅」、阪神沿線からは「尼崎駅」「武庫川駅」「出屋敷駅」へお越しいただき、市バスをご利用ください。どちらの駅からいらっしゃいますか？。１．ＪＲ立花駅から（位置的には、南西へ徒歩約１０分です。）。２．阪急塚口駅（南）から。３．阪急武庫之荘駅（南）から。４．阪神尼崎駅（北）から。５．阪神武庫川駅から。６．阪神出屋敷駅（北）から。＜改＞。【関連するＦＡＱ】。地域総合センターについて知りたい。＜改＞。【お問い合わせ】。地域総合センター今北。尼崎市西立花町３丁目１４－１。電話０６－６４１６－５７２９。
-----
Distance: 0.49538776305986054, QA_ID: 48, Answer: ■地域総合センター上ノ島本館及び分館には十分な駐車設備がございませんので、次のとおり阪神バス（尼崎市内線）をご利用ください。以下のいずれの駅からの行き方について知りたいですか。１．阪急塚口駅から。２．ＪＲ立花駅から。３．阪神尼崎駅から。＜改＞。【関連するＦＡＱ】。地域総合センターについて知りたい。＜改＞。【お問い合わせ】。地域総合センター上ノ島本館。尼崎市南塚口町８丁目７－２５。電話０６－６４２９－７６４０。
-----
Distance: 0.5507588593794948, QA_ID: 58, Answer: ■以下の場所で配布しています。【尾浜庁舎（道路維持担当窓口）】。【尼崎市役所本庁舎】。【尼崎市立魚つり公園】。【支所内地域振興センター】。【尼崎スポーツの森】。【尼崎の森中央緑地】。【尼崎総合庁舎】。【尼崎港管理事務所】。【北堀キャナルベース】。【尼ロック防災展示室】。【尼崎市内警察署】。【尼崎防犯協会】。【尼崎交通安全協会】。【消費生活センター】。【一般社団法人兵庫県自転車防犯登録会】。【ＪＲ尼崎駅案内所ｉ＋ＰｌｕｓＪＲ尼崎駅構内（改札外）】。【あまらぶアートラボＡ－ｌａｂ】。【中央公園内まちなかみどころご案内コーナー】。【阪神沿線レンタサイクル】。【尼崎市記念公園運動施設】。【スポーツクラブＷＯＯＤＹ】。【サンシビック尼崎】。【猪名川町スポーツセンター】。【市内の体育館】。各場所の詳細については次のリンクをご覧ください。［ＵＲＬ］。【お問い合わせ】。尼崎市コールセンター。電話０６－６３７５－５６３９。
-----
Distance: 0.5657978018642713, QA_ID: 33, Answer: 以下のいずれに該当されますか。１．引っ越し先が介護保険に関する施設以外の住宅。２．引っ越し先が介護保険に関する施設。３．受給資格証明書をもらっていない。＜改＞。■参考。介護保険事業担当でいう「介護保険に関する施設」とは、住所地特例対象施設のことをいい以下のとおりです。・介護保険３施設（特別養護老人ホーム、老人保健施設、介護療養型医療施設）。※ただし、地域密着型介護老人福祉施設（入所定員が３０人未満の特別養護老人ホーム）については、住所地特例の対象外となります。（１）特定施設（介護保険法第８条第１１項）。・養護老人ホーム（老人福祉法第２０条の４）。・軽費老人ホーム（老人福祉法第２０条の６）…ケアハウス等。・有料老人ホーム（老人福祉法第２９条第１項）。・有料老人ホームに該当するサービスを提供しているサービス付き高齢者向け住宅（高齢者の居住の安定確保に関する法律第５条第１項）。※特定施設は要介護認定の有無に関わらず、住所地特例の対象となります。ただし地域密着型特定施設（介護専用型で入居定員３０人未満）については住所地特例の対象外となります。＜改＞。詳しくは介護保険事業担当にお問い合わせください。■問合せ時間。午前８時４５分～午後５時３０分。ただし、窓口の取扱時間は午前９時～午後５時３０分。■休日。土・日曜日、祝日、年末年始（１２月２９日～１月３日）。■連絡先。【介護保険事業担当保険料担当】。電話０６－６４８９－６３７６。【北部保健福祉センター】。尼崎市南塚口町２－１－１。電話０６－４９５０－０５６２。【南部保健福祉センター】。尼崎市竹谷町２－１８３。電話０６－６４１５－６２７９。＜改＞。【関連するＦＡＱ】。尼崎市内から市外へ住所を移すときはどうしたらいいですか？（転出届）。
-----
Distance: 0.5857635336204101, QA_ID: 24, Answer: 以下のいずれに該当されますか。１．引っ越し先が市内の介護保険に関する施設以外の住宅。２．引っ越し先が市内の介護保険に関する施設。３．受給資格証明書をもらっていない。＜改＞。■参考。介護保険事業担当でいう「介護保険に関する施設」とは、住所地特例対象施設のことをいいます。・介護保険３施設（特別養護老人ホーム、老人保健施設、介護療養型医療施設）。※ただし、地域密着型介護老人福祉施設（入所定員が３０人未満の特別養護老人ホーム）については、住所地特例の対象外となります。・特定施設（介護保険法第８条第１１項）。（１）養護老人ホーム（老人福祉法第２０条の４）。（２）軽費老人ホーム（老人福祉法第２０条の６）…ケアハウス等。（３）有料老人ホーム（老人福祉法第２９条第１項）。（４）有料老人ホームに該当するサービスを提供しているサービス付き高齢者向け住宅（高齢者の居住の安定確保に関する法律第５条第１項）。※特定施設は要介護認定の有無に関わらず、住所地特例の対象となります。ただし地域密着型特定施設（介護専用型で入居定員３０人未満）については住所地特例の対象外となります。＜改＞。詳しくは介護保険事業担当へお問い合わせください。■問合せ時間：午前８時４５分～午後５時３０分。ただし、窓口の取扱時間は、午前９時～午後５時３０分。■休日：土・日曜日、祝日、年末年始（１２月２９日～１月３日）。■連絡先。【介護保険事業担当資格・保険料担当】。電話０６－６４８９－６３７６。【北部保健福祉センター】。尼崎市南塚口町２－１－１。電話０６－４９５０－０５６２。【南部保健福祉センター】。尼崎市竹谷町２－１８３。電話０６－６４１５－６２７９。【お問い合わせ】。健康福祉局福祉部。介護保険事業担当資格・保険料担当。電話０６－６４８９－６３７６。＜改＞。【関連するＦＡＱ】。市外から尼崎市内へ住所を移すときはどうしたらいいですか？（転入届）。
-----