[備忘録] Google Colabで30行！Qwen3-Embedding-0.6Bで日本語テキスト類似度計算 #Python – Qiita

# 必要なライブラリをインストール（Colabなら最初に一度だけ）
!pip install sentence-transformers torch -q

from sentence_transformers import SentenceTransformer, util

# モデル読み込み（GPUがあれば自動でGPUを使用）
model = SentenceTransformer("Qwen/Qwen3-Embedding-0.6B", trust_remote_code=True)

def compute_and_show(queries, documents, top_k=3):
    """類似度を計算して上位結果を表示"""
    print(f"🚀 {len(queries)}個のクエリ × {len(documents)}個の文書で類似度を計算中...")
    
    # 埋め込みを一括計算
    q_emb = model.encode(queries, convert_to_tensor=True)
    d_emb = model.encode(documents, convert_to_tensor=True)
    
    # コサイン類似度行列
    cos_scores = util.cos_sim(q_emb, d_emb)
    
    print(f"\n📊 各クエリの上位{top_k}件の結果:")
    print("="*60)
    
    # 各クエリの上位top_kを抽出＆表示
    for i, query in enumerate(queries):
        print(f"\n🔍 Q{i+1}: {query}")
        hits = util.semantic_search(q_emb[i:i+1], d_emb, top_k=top_k)[0]
        for rank, hit in enumerate(hits, 1):
            idx, score = hit['corpus_id'], hit['score']
            status = "✅" if rank == 1 and score > 0.7 else "📄"
            print(f"  {rank}位 {status} D{idx+1} ({score:.3f}): {documents[idx]}")

# 🍣 寿司をテーマにしたテスト
print("🍣 寿司をテーマにした類似度計算テスト")

queries = [
    "マグロの握り寿司について教えて",
    "回転寿司とは何ですか？", 
    "わさびの役割は？",
    "寿司の作り方を説明してください"
]

documents = [
    "マグロは寿司ネタの王様と呼ばれ、特に大トロは脂が多く濃厚な味わいが特徴です。",
    "回転寿司はベルトコンベアで寿司を提供するスタイルの飲食店です。",
    "わさびは魚の生臭さを消し、辛味で味を引き締める薬味です。",
    "寿司は酢飯の上に新鮮な魚介類を乗せた日本の伝統的な料理で、職人の技術が重要です。",
    "天ぷらは衣をつけて揚げる和食の代表料理です。"
]

# 実行
compute_and_show(queries, documents, top_k=3)

# 🏪 寿司屋さんでの会話シーン
print("\n" + "="*60)
print("🏪 寿司屋さんでの会話シーン")

customer_questions = [
    "今日のおすすめは何ですか？",
    "アレルギーがあるのですが大丈夫ですか？",
    "この魚はどこで獲れたものですか？"
]

shop_responses = [
    "本日のおすすめは朝獲れの天然本マグロの大トロです。脂の乗りが絶品です。",
    "えび、かに、卵などのアレルギー食材を使わないメニューもご用意しております。",
    "こちらの鯛は瀬戸内海で獲れた天然物で、身が締まって甘みがあります。",
    "当店では築地市場から毎朝仕入れた新鮮な魚介類のみを使用しています。"
]

compute_and_show(customer_questions, shop_responses, top_k=2)

print("\n🎉 実行完了！クエリと文書の類似度が正しく計算されました。")