［開発者向け］オンデバイスAIを加速する、軽量・高性能な新オープン埋め込みモデル「EmbeddingGemma」徹底解説

# 基本ライブラリのインストール
pip install -U sentence-transformers git+https://github.com/huggingface/transformers@v4.56.0-Embedding-Gemma-preview

# 追加で必要になる可能性があるライブラリ
pip install torch numpy datasets huggingface_hub

import torch
from sentence_transformers import SentenceTransformer
from huggingface_hub import login
import numpy as np

# Hugging Face Hubにログイン
login()  # トークンの入力が求められます

# デバイスの設定（GPU使用可能な場合はGPUを使用）
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"使用デバイス: {device}")

# EmbeddingGemmaモデルのロード
# 300Mパラメータ版を使用（軽量版）
model_id = "google/embeddinggemma-300M"
model = SentenceTransformer(model_id).to(device=device)

# モデル情報の表示
print(f"モデルデバイス: {model.device}")
print(f"総パラメータ数: {sum([p.numel() for _, p in model.named_parameters()]):,}")

# 埋め込みに変換したい文章のリスト
sentences = [
    "こんにちは、今日は良い天気ですね。",
    "最新のAI技術について調べています。",
    "猫は可愛い動物です。",
    "A cute animal is a cat.",
    "人工知能の研究が進んでいます。",
    "機械学習の勉強をしています。",
    "今日は雨が降っています。"
]

# タスク指定（推奨）：STS（Semantic Textual Similarity）を使用
task_name = "STS"

# 文章を埋め込みベクトルに変換
embeddings = model.encode(sentences, prompt=task_name)

print(f"埋め込みの形状: {embeddings.shape}")
print(f"ベクトル次元数: {embeddings.shape[1]}")

# 各文章の埋め込み結果を表示
for i, (sentence, embedding) in enumerate(zip(sentences, embeddings)):
    print(f"\n文章 {i+1}: {sentence}")
    print(f"埋め込み (最初の5次元): {embedding[:5]}")
    print(f"ベクトルの大きさ: {np.linalg.norm(embedding):.4f}")

def cosine_similarity(v1, v2):
    """コサイン類似度を計算する関数"""
    return np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))

def find_most_similar(query_embedding, document_embeddings, documents, top_k=3):
    """最も類似した文書を見つける関数"""
    similarities = []

    for doc_embedding in document_embeddings:
        similarity = cosine_similarity(query_embedding, doc_embedding)
        similarities.append(similarity)

    # 類似度でソート
    sorted_indices = np.argsort(similarities)[::-1]

    results = []
    for i in range(min(top_k, len(documents))):
        idx = sorted_indices[i]
        results.append({
            'document': documents[idx],
            'similarity': similarities[idx],
            'index': idx
        })

    return results

# 類似度計算の例
print("\n=== 類似度計算の例 ===")

# 日本語と英語の類似文の比較
japanese_cat = embeddings[2]  # "猫は可愛い動物です。"
english_cat = embeddings[3]   # "A cute animal is a cat."
similarity_cat = cosine_similarity(japanese_cat, english_cat)
print(f"日英の猫に関する文章の類似度: {similarity_cat:.4f}")

# AIに関する文章の比較
ai_research = embeddings[1]   # "最新のAI技術について調べています。"
ai_learning = embeddings[4]   # "人工知能の研究が進んでいます。"
similarity_ai = cosine_similarity(ai_research, ai_learning)
print(f"AI関連文章の類似度: {similarity_ai:.4f}")

# 機械学習の文章との比較
ml_study = embeddings[5]      # "機械学習の勉強をしています。"
similarity_ai_ml = cosine_similarity(ai_research, ml_study)
print(f"AI技術と機械学習の類似度: {similarity_ai_ml:.4f}")

def semantic_search(query, documents, model, task_name="STS", top_k=5):
    """セマンティック検索を実行する関数"""
    # クエリと文書の埋め込みを生成
    query_embedding = model.encode([query], prompt=task_name)[0]
    doc_embeddings = model.encode(documents, prompt=task_name)

    # 類似度計算
    similarities = []
    for doc_embedding in doc_embeddings:
        similarity = cosine_similarity(query_embedding, doc_embedding)
        similarities.append(similarity)

    # 結果をソート
    results = [(documents[i], similarities[i]) for i in range(len(documents))]
    results.sort(key=lambda x: x[1], reverse=True)

    return results[:top_k]

# セマンティック検索の例
knowledge_base = [
    "NISAアカウントの開設方法について説明します。",
    "普通預金口座の開設手続きをご案内します。",
    "住宅ローンの申し込みガイドをご確認ください。",
    "投資信託の管理手数料についてお知らせします。",
    "医療保険の補償内容をご説明します。",
    "生命保険の解約方針についてご案内します。",
    "税制優遇投資制度の利用方法を解説します。",
    "早期返済時の手数料について説明します。"
]

query = "税制優遇の投資を始めたいのですが、何をすればよいですか？"
print(f"\n=== セマンティック検索の例 ===")
print(f"クエリ: {query}")
print("\n検索結果:")

results = semantic_search(query, knowledge_base, model)
for i, (doc, score) in enumerate(results, 1):
    print(f"{i}. {doc} (類似度: {score:.4f})")

# MRLを使った異なる次元での埋め込み比較
def compare_embeddings_dimensions(sentences, model, dimensions=[128, 256, 512, 768]):
    """異なる次元での埋め込みの比較"""
    task_name = "STS"

    print("=== MRL次元比較 ===")

    for dim in dimensions:
        # 注意: 実際のMRL実装では、モデルの設定で出力次元を指定します
        # ここでは概念的な説明として、フル埋め込みから切り取ります
        embeddings = model.encode(sentences, prompt=task_name)
        truncated_embeddings = embeddings[:, :dim]

        print(f"\n次元数: {dim}")
        print(f"埋め込みサイズ: {truncated_embeddings.shape}")

        # 同じ文章間の類似度を計算（品質の指標として）
        if len(sentences) >= 2:
            sim = cosine_similarity(truncated_embeddings[0], truncated_embeddings[1])
            print(f"最初の2文の類似度: {sim:.4f}")

# MRL比較の実行
test_sentences = [
    "機械学習について学んでいます。",
    "人工知能の研究をしています。"
]

compare_embeddings_dimensions(test_sentences, model)

def batch_encode_with_progress(texts, model, batch_size=32, task_name="STS"):
    """大量のテキストを効率的にバッチ処理する関数"""
    embeddings = []

    print(f"総テキスト数: {len(texts)}")
    print(f"バッチサイズ: {batch_size}")

    for i in range(0, len(texts), batch_size):
        batch = texts[i:i+batch_size]
        batch_embeddings = model.encode(batch, prompt=task_name)
        embeddings.extend(batch_embeddings)

        print(f"処理済み: {min(i+batch_size, len(texts))}/{len(texts)}")

    return np.array(embeddings)

# 大量テキストの処理例
large_text_list = [
    f"これはテスト文書 {i} です。内容は様々な分野に関するものです。"
    for i in range(100)
]

# print("\n=== バッチ処理の例 ===")
# large_embeddings = batch_encode_with_progress(large_text_list, model, batch_size=16)

# EmbeddingGemmaで利用可能なタスク一覧
available_tasks = {
    "STS": "文の類似度計算（Semantic Textual Similarity）",
    "CLASSIFICATION": "テキスト分類",
    "CLUSTERING": "テキストクラスタリング",
    "RETRIEVAL_QUERY": "検索クエリ（検索用）",
    "RETRIEVAL_DOCUMENT": "検索文書（インデックス用）"
}

print("=== 利用可能なタスク ===")
for task, description in available_tasks.items():
    print(f"{task}: {description}")

# 異なるタスクでの埋め込み生成例
def compare_task_embeddings(text, model):
    """異なるタスクでの埋め込みを比較"""
    print(f"\nテキスト: {text}")

    for task in ["STS", "CLASSIFICATION", "CLUSTERING"]:
        try:
            embedding = model.encode([text], prompt=task)[0]
            print(f"{task}: 次元数 {len(embedding)}, ノルム {np.linalg.norm(embedding):.4f}")
        except Exception as e:
            print(f"{task}: エラー - {e}")

# タスク比較の実行
compare_task_embeddings("機械学習は人工知能の重要な分野です。", model)

# モデルの詳細情報を表示
def show_model_info(model):
    """モデルの詳細情報を表示する関数"""
    print("=== モデル詳細情報 ===")
    print(f"モデル: {model}")
    print(f"デバイス: {model.device}")

    # トークナイザー情報
    if hasattr(model, 'tokenizer'):
        print(f"最大シーケンス長: {model.max_seq_length}")
        print(f"トークナイザー: {type(model.tokenizer).__name__}")

    # 利用可能なプロンプト
    if hasattr(model, 'prompts'):
        print(f"利用可能なプロンプト: {list(model.prompts.keys())}")

    # パラメータ数の詳細
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

    print(f"総パラメータ数: {total_params:,}")
    print(f"学習可能パラメータ数: {trainable_params:,}")

# モデル情報の表示
show_model_info(model)

class DocumentSearchSystem:
    """EmbeddingGemmaを使用した文書検索システム"""

    def __init__(self, model):
        self.model = model
        self.documents = []
        self.embeddings = None

    def add_documents(self, documents):
        """文書を追加してインデックスを構築"""
        self.documents.extend(documents)
        new_embeddings = self.model.encode(documents, prompt="RETRIEVAL_DOCUMENT")

        if self.embeddings is None:
            self.embeddings = new_embeddings
        else:
            self.embeddings = np.vstack([self.embeddings, new_embeddings])

    def search(self, query, top_k=5):
        """クエリに対して検索を実行"""
        query_embedding = self.model.encode([query], prompt="RETRIEVAL_QUERY")[0]

        similarities = []
        for doc_embedding in self.embeddings:
            similarity = cosine_similarity(query_embedding, doc_embedding)
            similarities.append(similarity)

        # 上位結果を取得
        top_indices = np.argsort(similarities)[::-1][:top_k]

        results = []
        for idx in top_indices:
            results.append({
                'document': self.documents[idx],
                'similarity': similarities[idx],
                'index': idx
            })

        return results

# 使用例
search_system = DocumentSearchSystem(model)

# サンプル文書の追加
documents = [
    "EmbeddingGemmaは軽量で高性能なテキスト埋め込みモデルです。",
    "オンデバイスAIにより、プライバシーを保護しながら高度な検索が可能です。",
    "RAG（検索拡張生成）システムを構築できます。",
    "Matryoshka Representation Learningにより柔軟な次元調整が可能です。",
    "EdgeTPUで15ミリ秒未満の高速推論を実現します。"
]

search_system.add_documents(documents)

# 検索実行
query = "高速な推論について教えて"
results = search_system.search(query, top_k=3)

print(f"\nクエリ: {query}")
print("検索結果:")
for i, result in enumerate(results, 1):
    print(f"{i}. {result['document']} (類似度: {result['similarity']:.4f})")