最近話題のVector Searchを実現するFaissって何？ #1

pip install faiss-cpu
pip install sentence-transformers

import sqlite3
import numpy as np
from sentence_transformers import SentenceTransformer
import faiss

# 1. SQLiteデータベースの準備
def create_faq_database():
    conn = sqlite3.connect("faq.db")
    cursor = conn.cursor()
    cursor.execute("CREATE TABLE IF NOT EXISTS faq (id INTEGER PRIMARY KEY, question TEXT)")
    cursor.execute("DELETE FROM faq")

    # サンプルデータを追加
    sample_faqs = [
        "How do I reset my password?",
        "What payment methods do you accept?",
        "Can I return an item I bought?",
        "How can I track my order?"
    ]

    cursor.executemany("INSERT INTO faq (question) VALUES (?)", [(q,) for q in sample_faqs])
    conn.commit()
    return conn

# 2. 文章の埋め込みの作成
def compute_faq_embeddings(faq_questions, model):
    return np.array(model.encode(faq_questions))

# 3. Faissインデックスの作成
def create_faiss_index(model, embeddings, doc_ids):
    dimension = model.get_sentence_embedding_dimension()
    index_flat_l2 = faiss.IndexFlatL2(dimension)
    index = faiss.IndexIDMap(index_flat_l2)
    index.add_with_ids(embeddings, doc_ids)
    return index, index_flat_l2

# 4. クエリ処理
def search_faq(query, model, index, k=3):
    query_embedding = model.encode([query])
    distances, indices = index.search(np.array(query_embedding), k)
    return indices[0]

# 5. 結果の取得
def get_faq_results(faq_ids, conn):
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM faq WHERE id IN ({})".format(",".join("?" * len(faq_ids))), faq_ids)
    return cursor.fetchall()

# データベース作成
conn = create_faq_database()

# sentence-transformersモデルを読み込む
# model = SentenceTransformer("sentence-transformers/paraphrase-xlm-r-multilingual-v1")
model = SentenceTransformer("sentence-transformers/paraphrase-distilroberta-base-v1")

# FAQデータを読み込み、埋め込みベクトルを計算
cursor = conn.cursor()
cursor.execute("SELECT id, question FROM faq")
faq_data = cursor.fetchall()
faq_ids, faq_questions = zip(*faq_data)
faq_embeddings = compute_faq_embeddings(faq_questions, model)

# Faissインデックスを作成
index, index_flat_l2 = create_faiss_index(model, faq_embeddings, faq_ids)

# クエリを入力して検索
query = "How can I change my password?"
faq_indices = search_faq(query, model, index_flat_l2)

# 結果を取得して表示
results = get_faq_results([faq_ids[i] for i in faq_indices], conn)
print("Results for query:", query)
for r in results:
    print(f"ID: {r[0]}, Question: {r[1]}")

# データベース接続を閉じる
# conn.close()

def add_faq(question, conn, model, index):
    cursor = conn.cursor()
    cursor.execute("INSERT INTO faq (question) VALUES (?)", (question,))
    faq_id = cursor.lastrowid
    conn.commit()

    embedding = model.encode([question])
    index.add_with_ids(np.array(embedding), np.array([faq_id]))

    return faq_id

def remove_faq(faq_id, conn, index):
    cursor = conn.cursor()
    cursor.execute("DELETE FROM faq WHERE id=?", (faq_id,))
    conn.commit()

    # FaissインデックスからIDを削除
    id_selector = faiss.IDSelectorArray(np.array([faq_id], dtype=np.int64))
    index.remove_ids(id_selector)

def get_answer(query, model, index):
  faq_indices = search_faq(query, model, index)
  results = get_faq_results([faq_ids[i] for i in faq_indices], conn)
  return results

# FAQを追加
new_question = "What is your refund policy?"
new_faq_id = add_faq(new_question, conn, model, index)
print(f"Added new FAQ with ID: {new_faq_id}")

# 確認
query = "What's the refund polcy?"
faq_indices = search_faq(query, model, index)
print(f"{query}: {faq_indices}")

# FAQを削除
faq_id_to_remove = new_faq_id
remove_faq(faq_id_to_remove, conn, index)
print(f"Removed FAQ with ID: {faq_id_to_remove}")

# 確認
query = "What's the refund polcy?"
faq_indices = search_faq(query, model, index)
print(f"{query}: {faq_indices}")