埋め込みモデル完全ガイド: ベクトル検索・RAG・Sentence Transformers実践活用

# 直感的理解：単語をベクトルで表現
# "王" = [0.2, 0.8, 0.1, ...]
# "女王" = [0.3, 0.9, 0.1, ...]
# "男" = [0.1, 0.2, 0.8, ...]

# 有名なベクトル算術：king - man + woman ≈ queen
import numpy as np

king = np.array([0.2, 0.8, 0.1, 0.5])
man = np.array([0.1, 0.2, 0.8, 0.4])
woman = np.array([0.15, 0.25, 0.85, 0.6])
queen = np.array([0.3, 0.9, 0.1, 0.7])

result = king - man + woman
print(f"king - man + woman = {result}")
print(f"queen              = {queen}")
# 2つのベクトルが非常に類似していることが確認できる

# モデル選択の意思決定ツリー
def select_embedding_model(requirements):
    """要件に基づく埋め込みモデル選択ガイド"""

    if requirements.get("budget") == "unlimited":
        if requirements.get("max_performance"):
            return "GTE-Qwen2-7B-instruct（セルフホスト）またはVoyage-3（API）"
        return "text-embedding-3-large（OpenAI API）"

    if requirements.get("multilingual"):
        if requirements.get("self_hosted"):
            return "BGE-M3（Dense+Sparseハイブリッド）"
        return "Cohere embed-v3 multilingual"

    if requirements.get("low_latency"):
        if requirements.get("self_hosted"):
            return "nomic-embed-text-v1.5（軽量137M）"
        return "text-embedding-3-small（OpenAI API）"

    if requirements.get("domain_specific"):
        return "Sentence Transformers + ファインチューニング（ベースモデル：BGEまたはE5）"

    # デフォルト推奨
    return "text-embedding-3-small（コスト効率の良い汎用選択）"

from sentence_transformers import SentenceTransformer
import numpy as np

# モデルのロード
model = SentenceTransformer('BAAI/bge-large-en-v1.5')

# 単一文の埋め込み
sentence = "Embedding models convert text into numerical vectors."
embedding = model.encode(sentence)
print(f"次元: {embedding.shape}")  # (1024,)

# バッチ埋め込み
sentences = [
    "埋め込みモデルはテキストを数値ベクトルに変換する。",
    "ベクトル検索は類似文書を高速に検索する。",
    "RAGは検索ベースの生成技法である。",
    "今日の天気はとても良い。",
]

embeddings = model.encode(sentences, batch_size=32, show_progress_bar=True)
print(f"埋め込み行列のサイズ: {embeddings.shape}")  # (4, 1024)

# 類似度計算
from sentence_transformers.util import cos_sim

similarity_matrix = cos_sim(embeddings, embeddings)
print("類似度行列:")
print(similarity_matrix.numpy().round(3))

from sentence_transformers import SentenceTransformer

# BGE-M3: 100以上の言語をサポートする多言語埋め込みモデル
model = SentenceTransformer('BAAI/bge-m3')

# 多言語文埋め込み
sentences = [
    "Machine learning is transforming the world.",        # 英語
    "머신러닝이 세상을 변화시키고 있다.",                      # 韓国語
    "機械学習が世界を変えている。",                           # 日本語
    "机器学习正在改变世界。",                                # 中国語
]

embeddings = model.encode(sentences, normalize_embeddings=True)

# 言語間類似度の確認
from sentence_transformers.util import cos_sim

similarities = cos_sim(embeddings, embeddings)
print("多言語類似度行列:")
for i, s1 in enumerate(sentences):
    for j, s2 in enumerate(sentences):
        if i < j:
            print(f"  '{s1[:20]}...' <-> '{s2[:20]}...': {similarities[i][j]:.4f}")
# 同じ意味の異なる言語の文が高い類似度を示す

from openai import OpenAI
import numpy as np

client = OpenAI()

def get_openai_embeddings(texts, model="text-embedding-3-small", dimensions=None):
    """OpenAI埋め込み生成（Matryoshka次元削減対応）"""
    kwargs = {"input": texts, "model": model}
    if dimensions:
        kwargs["dimensions"] = dimensions

    response = client.embeddings.create(**kwargs)
    return np.array([item.embedding for item in response.data])

# 基本的な使用方法
texts = ["埋め込みモデルの原理", "ベクトル検索システムの構築"]
embeddings_full = get_openai_embeddings(texts, model="text-embedding-3-large")
print(f"全次元: {embeddings_full.shape}")  # (2, 3072)

# Matryoshka: 次元削減によるコスト・速度最適化
embeddings_256 = get_openai_embeddings(
    texts, model="text-embedding-3-large", dimensions=256
)
print(f"削減次元: {embeddings_256.shape}")  # (2, 256)

# コサイン類似度の比較
def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

sim_full = cosine_similarity(embeddings_full[0], embeddings_full[1])
sim_256 = cosine_similarity(embeddings_256[0], embeddings_256[1])
print(f"全次元類似度: {sim_full:.4f}")
print(f"256次元類似度: {sim_256:.4f}")

import faiss
import numpy as np
import time

# テストデータの生成
np.random.seed(42)
dimension = 1024
num_vectors = 1_000_000
num_queries = 100

# 正規化されたランダムベクトルの生成
data = np.random.randn(num_vectors, dimension).astype('float32')
faiss.normalize_L2(data)
queries = np.random.randn(num_queries, dimension).astype('float32')
faiss.normalize_L2(queries)

# 1. Flat Index（正確検索、Brute-force）
print("=== Flat Index（正確検索） ===")
index_flat = faiss.IndexFlatIP(dimension)  # Inner Product（コサイン類似度）
index_flat.add(data)

start = time.time()
D_exact, I_exact = index_flat.search(queries, 10)
flat_time = time.time() - start
print(f"検索時間: {flat_time:.3f}秒")
print(f"Recall@10: 1.000（正確検索）")

# 2. IVF（Inverted File Index）
print("\n=== IVF Index ===")
nlist = 1024  # クラスタ数
quantizer = faiss.IndexFlatIP(dimension)
index_ivf = faiss.IndexIVFFlat(quantizer, dimension, nlist, faiss.METRIC_INNER_PRODUCT)
index_ivf.train(data)
index_ivf.add(data)
index_ivf.nprobe = 32  # 検索するクラスタ数

start = time.time()
D_ivf, I_ivf = index_ivf.search(queries, 10)
ivf_time = time.time() - start

# Recall計算
recall = np.mean([len(set(I_ivf[i]) & set(I_exact[i])) / 10 for i in range(num_queries)])
print(f"検索時間: {ivf_time:.3f}秒（x{flat_time/ivf_time:.1f}高速）")
print(f"Recall@10: {recall:.3f}")

# 3. HNSW（Hierarchical Navigable Small World）
print("\n=== HNSW Index ===")
index_hnsw = faiss.IndexHNSWFlat(dimension, 32)  # M=32
index_hnsw.hnsw.efConstruction = 200
index_hnsw.hnsw.efSearch = 64
index_hnsw.metric_type = faiss.METRIC_INNER_PRODUCT
index_hnsw.add(data)

start = time.time()
D_hnsw, I_hnsw = index_hnsw.search(queries, 10)
hnsw_time = time.time() - start

recall = np.mean([len(set(I_hnsw[i]) & set(I_exact[i])) / 10 for i in range(num_queries)])
print(f"検索時間: {hnsw_time:.3f}秒（x{flat_time/hnsw_time:.1f}高速）")
print(f"Recall@10: {recall:.3f}")

# 4. IVF-PQ（Product Quantization）
print("\n=== IVF-PQ Index（メモリ最適化） ===")
m = 64  # サブベクトル数
nbits = 8  # コードブックビット数
index_ivfpq = faiss.IndexIVFPQ(quantizer, dimension, nlist, m, nbits)
index_ivfpq.train(data)
index_ivfpq.add(data)
index_ivfpq.nprobe = 32

start = time.time()
D_pq, I_pq = index_ivfpq.search(queries, 10)
pq_time = time.time() - start

recall = np.mean([len(set(I_pq[i]) & set(I_exact[i])) / 10 for i in range(num_queries)])
print(f"検索時間: {pq_time:.3f}秒（x{flat_time/pq_time:.1f}高速）")
print(f"Recall@10: {recall:.3f}")
print(f"メモリ: Flat={data.nbytes/1e9:.1f}GB, PQ={index_ivfpq.sa_code_size()*num_vectors/1e9:.3f}GB")

import chromadb
from chromadb.utils import embedding_functions

# Chromaクライアントの作成
client = chromadb.PersistentClient(path="./chroma_db")

# Sentence Transformers埋め込み関数の設定
embedding_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="BAAI/bge-m3"
)

# コレクションの作成
collection = client.get_or_create_collection(
    name="tech_documents",
    embedding_function=embedding_fn,
    metadata={"hnsw:space": "cosine"}  # コサイン類似度を使用
)

# ドキュメントの追加
documents = [
    "Kubernetesはコンテナオーケストレーションプラットフォームである。",
    "Dockerはアプリケーションをコンテナにパッケージングするツールである。",
    "Prometheusはメトリクスベースのモニタリングシステムである。",
    "Grafanaはデータ可視化およびダッシュボードツールである。",
    "TerraformはインフラをコードとしてIaCで管理するツールである。",
    "埋め込みモデルはテキストをベクトルに変換する。",
    "RAGは検索拡張生成技法でLLMのハルシネーションを削減する。",
]

collection.add(
    documents=documents,
    ids=[f"doc_{i}" for i in range(len(documents))],
    metadatas=[
        {"category": "kubernetes"},
        {"category": "docker"},
        {"category": "monitoring"},
        {"category": "monitoring"},
        {"category": "iac"},
        {"category": "ai"},
        {"category": "ai"},
    ]
)

# セマンティック検索
results = collection.query(
    query_texts=["コンテナ関連の技術は？"],
    n_results=3
)
print("検索結果:")
for doc, score in zip(results["documents"][0], results["distances"][0]):
    print(f"  [{score:.4f}] {doc}")

# メタデータフィルタリング + セマンティック検索
results_filtered = collection.query(
    query_texts=["モニタリングツール"],
    n_results=2,
    where={"category": "monitoring"}
)
print("\nフィルタリング済み検索結果:")
for doc in results_filtered["documents"][0]:
    print(f"  {doc}")

import numpy as np

def cosine_similarity(a, b):
    """コサイン類似度：ベクトル方向の類似性を測定"""
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

def dot_product(a, b):
    """内積：正規化されたベクトルではコサイン類似度と同等"""
    return np.dot(a, b)

def euclidean_distance(a, b):
    """ユークリッド距離：ベクトル間の直線距離"""
    return np.linalg.norm(a - b)

def manhattan_distance(a, b):
    """マンハッタン距離：次元ごとの絶対差の合計"""
    return np.sum(np.abs(a - b))

# 例示ベクトル
a = np.array([1.0, 2.0, 3.0])
b = np.array([1.0, 2.0, 3.1])
c = np.array([-1.0, -2.0, -3.0])

print("ベクトルaとb（非常に類似）:")
print(f"  コサイン類似度:    {cosine_similarity(a, b):.4f}")
print(f"  ユークリッド距離:  {euclidean_distance(a, b):.4f}")
print(f"  内積:             {dot_product(a, b):.4f}")

print("\nベクトルaとc（反対方向）:")
print(f"  コサイン類似度:    {cosine_similarity(a, c):.4f}")
print(f"  ユークリッド距離:  {euclidean_distance(a, c):.4f}")
print(f"  内積:             {dot_product(a, c):.4f}")

from sentence_transformers import SentenceTransformer, util
import torch

class SemanticSearchEngine:
    def __init__(self, model_name="BAAI/bge-m3"):
        self.model = SentenceTransformer(model_name)
        self.documents = []
        self.embeddings = None

    def index_documents(self, documents):
        """ドキュメントをインデキシングして埋め込みを生成"""
        self.documents = documents
        self.embeddings = self.model.encode(
            documents,
            convert_to_tensor=True,
            normalize_embeddings=True,
            show_progress_bar=True
        )
        print(f"{len(documents)}件のドキュメントをインデキシング完了（次元: {self.embeddings.shape[1]}）")

    def search(self, query, top_k=5):
        """セマンティック検索の実行"""
        query_embedding = self.model.encode(
            query,
            convert_to_tensor=True,
            normalize_embeddings=True
        )

        # コサイン類似度の計算
        scores = util.cos_sim(query_embedding, self.embeddings)[0]

        # 上位k件の結果を返す
        top_results = torch.topk(scores, k=min(top_k, len(self.documents)))

        results = []
        for score, idx in zip(top_results.values, top_results.indices):
            results.append({
                "document": self.documents[idx],
                "score": score.item(),
                "index": idx.item()
            })
        return results

    def search_with_reranking(self, query, top_k=5, initial_k=20):
        """2段階検索：埋め込み検索 + リランキング"""
        from sentence_transformers import CrossEncoder

        # 第1段階：埋め込みベースの候補検索
        candidates = self.search(query, top_k=initial_k)

        # 第2段階：Cross-encoderでリランキング
        reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-12-v2')
        pairs = [(query, c["document"]) for c in candidates]
        rerank_scores = reranker.predict(pairs)

        # リランキング結果を返す
        for i, score in enumerate(rerank_scores):
            candidates[i]["rerank_score"] = float(score)

        reranked = sorted(candidates, key=lambda x: x["rerank_score"], reverse=True)
        return reranked[:top_k]

# 使用例
engine = SemanticSearchEngine()

documents = [
    "Pythonはデータサイエンスと機械学習で最も使われているプログラミング言語である。",
    "JavaScriptはWeb開発の中核言語で、Node.jsによりサーバーサイドでも使われる。",
    "Kubernetesはコンテナ化されたアプリケーションのデプロイ、スケーリング、管理を自動化する。",
    "PostgreSQLは強力なオープンソースのリレーショナルデータベース管理システムである。",
    "TensorFlowとPyTorchはディープラーニングモデル開発で最も広く使われるフレームワークである。",
    "Redisはインメモリデータ構造ストアで、キャッシュとメッセージブローカーとして活用される。",
    "Dockerはアプリケーションと依存関係をコンテナにパッケージングしてポータビリティを提供する。",
    "GraphQLはRESTの代替として、クライアントが必要なデータのみをリクエストできるようにする。",
]

engine.index_documents(documents)

# セマンティック検索
query = "ディープラーニング開発にはどのツールを使うべきですか？"
results = engine.search(query, top_k=3)
print(f"\nクエリ: {query}")
for r in results:
    print(f"  [{r['score']:.4f}] {r['document']}")

from sentence_transformers import SentenceTransformer, CrossEncoder
from openai import OpenAI
import chromadb
from chromadb.utils import embedding_functions
import tiktoken
from typing import List, Dict

class RAGPipeline:
    def __init__(
        self,
        embedding_model: str = "BAAI/bge-m3",
        reranker_model: str = "cross-encoder/ms-marco-MiniLM-L-12-v2",
        llm_model: str = "gpt-4o",
    ):
        self.embedder = SentenceTransformer(embedding_model)
        self.reranker = CrossEncoder(reranker_model)
        self.llm_client = OpenAI()
        self.llm_model = llm_model

        # Chromaベクトルデータベースの初期化
        self.chroma_client = chromadb.PersistentClient(path="./rag_db")
        self.collection = self.chroma_client.get_or_create_collection(
            name="rag_documents",
            metadata={"hnsw:space": "cosine"}
        )

    def chunk_text(self, text: str, chunk_size: int = 512, overlap: int = 50) -> List[str]:
        """トークンベースのテキストチャンク分割"""
        tokenizer = tiktoken.get_encoding("cl100k_base")
        tokens = tokenizer.encode(text)
        chunks = []

        start = 0
        while start < len(tokens):
            end = start + chunk_size
            chunk_tokens = tokens[start:end]
            chunk_text = tokenizer.decode(chunk_tokens)
            chunks.append(chunk_text)
            start = end - overlap  # オーバーラップの適用

        return chunks

    def ingest_documents(self, documents: List[Dict[str, str]]):
        """ドキュメントをチャンク分割してベクトルDBに保存"""
        all_chunks = []
        all_ids = []
        all_metadatas = []

        for doc_idx, doc in enumerate(documents):
            chunks = self.chunk_text(doc["content"])
            for chunk_idx, chunk in enumerate(chunks):
                all_chunks.append(chunk)
                all_ids.append(f"doc{doc_idx}_chunk{chunk_idx}")
                all_metadatas.append({
                    "source": doc.get("source", "unknown"),
                    "doc_index": doc_idx,
                    "chunk_index": chunk_idx,
                })

        # 埋め込み生成および保存
        embeddings = self.embedder.encode(all_chunks, normalize_embeddings=True)

        self.collection.add(
            documents=all_chunks,
            embeddings=embeddings.tolist(),
            ids=all_ids,
            metadatas=all_metadatas,
        )
        print(f"{len(documents)}件のドキュメント -> {len(all_chunks)}件のチャンクをインデキシング完了")

    def retrieve(self, query: str, top_k: int = 10) -> List[Dict]:
        """ベクトル類似度ベースの検索"""
        query_embedding = self.embedder.encode(
            [query], normalize_embeddings=True
        ).tolist()

        results = self.collection.query(
            query_embeddings=query_embedding,
            n_results=top_k,
        )

        retrieved = []
        for i in range(len(results["documents"][0])):
            retrieved.append({
                "text": results["documents"][0][i],
                "metadata": results["metadatas"][0][i],
                "distance": results["distances"][0][i],
            })
        return retrieved

    def rerank(self, query: str, candidates: List[Dict], top_k: int = 5) -> List[Dict]:
        """Cross-encoderベースのリランキング"""
        pairs = [(query, c["text"]) for c in candidates]
        scores = self.reranker.predict(pairs)

        for i, score in enumerate(scores):
            candidates[i]["rerank_score"] = float(score)

        reranked = sorted(candidates, key=lambda x: x["rerank_score"], reverse=True)
        return reranked[:top_k]

    def generate(self, query: str, context_docs: List[Dict]) -> str:
        """検索されたコンテキストに基づいてLLM応答を生成"""
        context = "\n\n---\n\n".join([doc["text"] for doc in context_docs])

        messages = [
            {
                "role": "system",
                "content": (
                    "You are a helpful assistant. Answer the question based on "
                    "the provided context. If the context doesn't contain "
                    "relevant information, say so."
                ),
            },
            {
                "role": "user",
                "content": f"Context:\n{context}\n\nQuestion: {query}",
            },
        ]

        response = self.llm_client.chat.completions.create(
            model=self.llm_model,
            messages=messages,
            temperature=0.1,
        )
        return response.choices[0].message.content

    def query(self, question: str, top_k_retrieve: int = 10, top_k_rerank: int = 5) -> str:
        """RAGパイプライン全体の実行"""
        # ステップ1：検索
        candidates = self.retrieve(question, top_k=top_k_retrieve)
        print(f"第1段階検索: {len(candidates)}件の候補")

        # ステップ2：リランキング
        reranked = self.rerank(question, candidates, top_k=top_k_rerank)
        print(f"第2段階リランキング: 上位{len(reranked)}件を選択")

        # ステップ3：生成
        answer = self.generate(question, reranked)
        return answer

# 使用例
rag = RAGPipeline()

# ドキュメントのインジェスト
documents = [
    {"content": "長い技術ドキュメントの内容...", "source": "tech_doc_1.pdf"},
    {"content": "別のドキュメントの内容...", "source": "tech_doc_2.pdf"},
]
rag.ingest_documents(documents)

# クエリ
answer = rag.query("埋め込みモデルの次元数は性能にどのような影響を与えますか？")
print(f"\n回答: {answer}")

from rank_bm25 import BM25Okapi
import numpy as np

class HybridSearchEngine:
    """Dense（埋め込み）+ Sparse（BM25）ハイブリッド検索"""

    def __init__(self, embedding_model="BAAI/bge-m3"):
        self.embedder = SentenceTransformer(embedding_model)
        self.documents = []
        self.embeddings = None
        self.bm25 = None

    def index(self, documents):
        self.documents = documents

        # Dense：埋め込み生成
        self.embeddings = self.embedder.encode(
            documents, normalize_embeddings=True, convert_to_tensor=True
        )

        # Sparse：BM25インデックス構築
        tokenized = [doc.split() for doc in documents]
        self.bm25 = BM25Okapi(tokenized)

    def search(self, query, top_k=5, alpha=0.7):
        """ハイブリッド検索（alpha: dense重み、1-alpha: sparse重み）"""
        # Dense検索
        query_emb = self.embedder.encode(
            query, normalize_embeddings=True, convert_to_tensor=True
        )
        dense_scores = util.cos_sim(query_emb, self.embeddings)[0].cpu().numpy()

        # Sparse検索（BM25）
        sparse_scores = self.bm25.get_scores(query.split())

        # 正規化
        if dense_scores.max() > 0:
            dense_scores = dense_scores / dense_scores.max()
        if sparse_scores.max() > 0:
            sparse_scores = sparse_scores / sparse_scores.max()

        # 加重合算
        hybrid_scores = alpha * dense_scores + (1 - alpha) * sparse_scores

        # 上位k件を返す
        top_indices = np.argsort(hybrid_scores)[::-1][:top_k]
        return [
            {
                "document": self.documents[i],
                "score": hybrid_scores[i],
                "dense_score": dense_scores[i],
                "sparse_score": sparse_scores[i],
            }
            for i in top_indices
        ]

from sentence_transformers import (
    SentenceTransformer,
    InputExample,
    losses,
    evaluation,
)
from torch.utils.data import DataLoader

# ベースモデルのロード
model = SentenceTransformer('BAAI/bge-base-en-v1.5')

# 学習データの準備（anchor、positive、negative）
train_examples = [
    # （クエリ、関連ドキュメント、非関連ドキュメント）
    InputExample(texts=[
        "How to deploy a Kubernetes pod?",
        "kubectl apply -f pod.yaml creates a new pod in the cluster.",
        "Python is a popular programming language for data science."
    ]),
    InputExample(texts=[
        "What is a Docker container?",
        "A Docker container is a lightweight, standalone executable package.",
        "Machine learning models require large datasets for training."
    ]),
    InputExample(texts=[
        "How does Redis caching work?",
        "Redis stores data in memory for fast read/write access as a cache layer.",
        "Kubernetes orchestrates containerized applications across clusters."
    ]),
    # ... 数千〜数万件の学習例
]

# DataLoaderの作成
train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=16)

# 損失関数：TripletLoss（anchor, positive, negative）
train_loss = losses.TripletLoss(model=model)

# 評価データ
eval_examples = [
    InputExample(texts=["query1", "relevant_doc1"], label=1.0),
    InputExample(texts=["query2", "irrelevant_doc2"], label=0.0),
]
evaluator = evaluation.EmbeddingSimilarityEvaluator.from_input_examples(
    eval_examples, name="domain-eval"
)

# ファインチューニングの実行
model.fit(
    train_objectives=[(train_dataloader, train_loss)],
    evaluator=evaluator,
    epochs=3,
    warmup_steps=100,
    evaluation_steps=500,
    output_path="./finetuned_embedding_model",
    save_best_model=True,
)

print("ファインチューニング完了！")

# ファインチューニング済みモデルのロードと使用
finetuned_model = SentenceTransformer('./finetuned_embedding_model')
embeddings = finetuned_model.encode(["ドメイン特化クエリ"])

from sentence_transformers import SentenceTransformer, InputExample, losses
from torch.utils.data import DataLoader

model = SentenceTransformer('BAAI/bge-base-en-v1.5')

# (query, positive_passage)ペアのみで学習可能
# in-batch negativesを自動的に活用
train_examples = [
    InputExample(texts=["What is embedding?", "An embedding is a vector representation of data."]),
    InputExample(texts=["How does HNSW work?", "HNSW builds a hierarchical graph for approximate nearest neighbor search."]),
    InputExample(texts=["What is RAG?", "RAG retrieves relevant documents and uses them to augment LLM generation."]),
    # ... さらに多くの(query, positive)ペア
]

train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=32)

# MultipleNegativesRankingLoss: バッチ内の他のpositiveをnegativeとして活用
train_loss = losses.MultipleNegativesRankingLoss(model=model)

model.fit(
    train_objectives=[(train_dataloader, train_loss)],
    epochs=3,
    warmup_steps=100,
    output_path="./mnrl_finetuned_model",
)

from sentence_transformers import SentenceTransformer
import numpy as np

# Matryoshka Representation Learning（MRL）対応モデル
model = SentenceTransformer('nomic-ai/nomic-embed-text-v1.5')

texts = [
    "Vector databases store embeddings for similarity search.",
    "Embedding models convert text into numerical representations.",
    "RAG systems combine retrieval with language generation.",
]

# 全次元の埋め込み
full_embeddings = model.encode(texts)
print(f"全次元: {full_embeddings.shape[1]}")  # 768

# Matryoshka：希望の次元に切り詰めて正規化
def truncate_embeddings(embeddings, target_dim):
    """Matryoshka方式による次元削減"""
    truncated = embeddings[:, :target_dim]
    # L2正規化
    norms = np.linalg.norm(truncated, axis=1, keepdims=True)
    return truncated / norms

# 様々な次元での類似度比較
for dim in [64, 128, 256, 512, 768]:
    reduced = truncate_embeddings(full_embeddings, dim)
    sim = np.dot(reduced[0], reduced[1])  # 正規化ベクトルの内積 = コサイン類似度
    print(f"  次元 {dim:>4}: 類似度 = {sim:.4f}")

import numpy as np

def scalar_quantize_int8(embeddings):
    """スカラー量子化：float32 -> int8（メモリ75%削減）"""
    min_val = embeddings.min(axis=0)
    max_val = embeddings.max(axis=0)
    scale = (max_val - min_val) / 255.0

    quantized = ((embeddings - min_val) / scale).astype(np.int8)
    return quantized, min_val, scale

def scalar_dequantize_int8(quantized, min_val, scale):
    """逆量子化：int8 -> float32"""
    return quantized.astype(np.float32) * scale + min_val

def binary_quantize(embeddings):
    """バイナリ量子化：float32 -> 1bit（メモリ32倍削減）"""
    return (embeddings > 0).astype(np.uint8)

# メモリ比較
num_vectors = 1_000_000
dimension = 1024
embeddings = np.random.randn(num_vectors, dimension).astype(np.float32)

print(f"オリジナル（float32）: {embeddings.nbytes / 1e9:.2f} GB")

quantized, _, _ = scalar_quantize_int8(embeddings)
print(f"int8量子化:           {quantized.nbytes / 1e9:.2f} GB")

binary = binary_quantize(embeddings)
print(f"バイナリ量子化:       {binary.nbytes / 1e9:.2f} GB")
# オリジナル（float32）: 4.10 GB
# int8量子化:           1.02 GB
# バイナリ量子化:       1.02 GB（実際のビットパッキング時0.13 GB）

import asyncio
from sentence_transformers import SentenceTransformer
from functools import lru_cache
import hashlib

class OptimizedEmbeddingService:
    def __init__(self, model_name="BAAI/bge-m3", cache_size=10000):
        self.model = SentenceTransformer(model_name)
        self.cache = {}
        self.cache_size = cache_size

    def _get_cache_key(self, text):
        return hashlib.md5(text.encode()).hexdigest()

    def encode_with_cache(self, texts, batch_size=64):
        """キャッシュを適用した埋め込み生成"""
        uncached_texts = []
        uncached_indices = []
        results = [None] * len(texts)

        # キャッシュヒットの確認
        for i, text in enumerate(texts):
            key = self._get_cache_key(text)
            if key in self.cache:
                results[i] = self.cache[key]
            else:
                uncached_texts.append(text)
                uncached_indices.append(i)

        # キャッシュミスのテキストのみバッチ埋め込み
        if uncached_texts:
            new_embeddings = self.model.encode(
                uncached_texts,
                batch_size=batch_size,
                normalize_embeddings=True,
            )

            for idx, emb in zip(uncached_indices, new_embeddings):
                key = self._get_cache_key(texts[idx])
                self.cache[key] = emb
                results[idx] = emb

                # キャッシュサイズの管理
                if len(self.cache) > self.cache_size:
                    oldest_key = next(iter(self.cache))
                    del self.cache[oldest_key]

        return results

    def get_cache_stats(self):
        return {"cache_size": len(self.cache), "max_size": self.cache_size}