FAISS

src/FAISS/faiss.py

import faiss
import openai
import numpy as np
import os

class FAISSWrapper:

    def __init__(self, embedding_model: str = "text-embedding-3-small"):
        """
        コンストラクタ。OpenAI APIキーと埋め込みモデルを初期化します。
        """
        self.embedding_model = embedding_model
        self.index = None  # FAISSインデックスを格納
        self.corpus = []  # オリジナルのテキストをキャッシュ

    def _get_embeddings(self, texts: list[str]) -> np.ndarray:
        """
        OpenAIの埋め込みモデルを使用してテキストをベクトル化します。
        """
        client = openai.AzureOpenAI(
            azure_endpoint=os.environ.get("AZURE_OPENAI_ENDPOINT"),
            api_key=os.environ.get("AZURE_OPENAI_API_KEY"),
            api_version=os.environ.get("OPENAI_API_VERSION")
        )
        response = client.embeddings.create(
            model=self.embedding_model,  # OpenAIの埋め込みモデル
            input=texts
        )
        
        embeddings = [d.embedding for d in response.data]
        return np.array(embeddings, dtype='float32')

    def embedding_corpus(self, corpus: list[str]) -> np.ndarray:
        """
        コーパス全体をベクトル化し、FAISSインデックスを構築します。
        """
        self.corpus = corpus
        embeddings = self._get_embeddings(corpus)

        # ベクトル次元を取得し、FAISSインデックスを初期化
        dimension = embeddings.shape[1]
        self.index = faiss.IndexFlatL2(dimension)
        self.index.add(embeddings)  # インデックスにベクトルを追加
        return embeddings

    def rank_documents(self, query: str, top_k: int = 10) -> tuple[np.ndarray, np.ndarray]:
        """
        クエリに対して類似度の高い文書をランキングする。
        """
        if self.index is None:
            raise ValueError("FAISSインデックスが初期化されていません。embedding_corpus()を呼び出してください。")
        
        # クエリをベクトル化
        query_embedding = self._get_embeddings([query])

        # 類似度検索を実行
        distances, indices = self.index.search(query_embedding, top_k)
        return indices[0], distances[0]

    def get_original_texts(self, indices: list[int]) -> list[str]:
        """
        インデックスからオリジナルのテキストを取得する。
        """
        return [self.corpus[i] for i in indices]

# 使用例
if __name__ == "__main__":
    # 環境変数からOpenAI APIキーを取得（手動で設定も可）
    OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "YOUR_API_KEY_HERE")

    # テキストコーパス
    texts = [
        "This is a test sentence.",
        "I love programming.",
        "Artificial Intelligence is fascinating.",
        "FAISS is a library for similarity search.",
        "Python is great for data science."
    ]

    # クラスの初期化
    faiss_wrapper = FAISSWrapper(openai_api_key=OPENAI_API_KEY)

    # コーパスをベクトル化してFAISSに登録
    faiss_wrapper.embedding_corpus(texts)

    # クエリを実行
    query = "I like machine learning."
    distances, indices = faiss_wrapper.rank_documents(query)

    # 結果を表示
    print("\nクエリ:", query)
    print("\n類似したテキスト:")
    for i, idx in enumerate(indices):
        print(f"{i+1}: {faiss_wrapper.corpus[idx]} (距離: {distances[i]:.4f})")