Bedrock Cohere Embed の日本語性能を Titan Embeddings と比較する - Amazon Titan Embeddingsでベクトル入門⑤

Cohere EmbedなるEmbeddingモデルが使用可能になったので、Amazon Titan Embeddingsと比較してみます。

比較対象は以下

作ったもの

まずはガワを作って比較してみます。2023/11/14現在、LangChainのBedrockEmbeddingsでは上手く動かなかった（パラメータとレスポンスがTitan Embeddings用になってる？）のでboto3で書いてます。

よく考えると'texts'に1も2も入れてまとめてEmbedding出来る気がしますがまあこのまま…

バージョン0.0.336でlangChainがBedrockのCohere Embedに対応したのでLangChain版に変更しました。

cos2_cohere_embed.py

from langchain.embeddings import BedrockEmbeddings
from sklearn.metrics.pairwise import cosine_similarity
import streamlit as st

st.title("コサイン類似度 測ルンです  \n（Cohere Embed）")
input_text_1 = st.text_area("文章1")
input_text_2 = st.text_area("文章2")
send_button = st.button("送信")

if send_button:

    # 文章1,2のEmbedding
    embeddings = BedrockEmbeddings(model_id="cohere.embed-multilingual-v3")
    vector_1 = embeddings.embed_documents([input_text_1])
    vector_2 = embeddings.embed_documents([input_text_2])

    # コサイン類似度の計算
    similarities = cosine_similarity(vector_1, vector_2).tolist()[0][0]
    st.write("コサイン類似度: " + str(similarities))    

    st.text("文章1のVector: ")
    st.write(vector_1)    
    st.text("文章2のVector: ")
    st.write(vector_2)   

実行

python -m streamlit run cos2_cohere_embed.py

文章1に入力されたものと文章2に入力されたもののコサイン類似度を出力するプログラムです。

以下のデータセット（JGLUE:日本語言語理解ベンチマークのJSTSデータセット）から幾つかピックアップして、sentence1とsentence2のコサイン類似度を見てみます。

labelが0.0のものを幾つか試します。

コサイン類似度: 0.5316826458985118

コサイン類似度: 0.41077689352293467

コサイン類似度: 0.561588540325247

こちらも関係なさそうな文章のコサイン類似度は0.5程度になりました。
今度はlabelが4.xのものを幾つか試します。

コサイン類似度: 0.7746420939787051

コサイン類似度: 0.9753954063523695

コサイン類似度: 0.8917781565650358

こちらも明らかに高い値になりました。Titan Embeddingsと似たような値になりますね。
最後にlabelが3.0のものを試してみます。

コサイン類似度: 0.8770156997178417

これもTitan Embeddingsと同じ傾向で、上記のlabelが4.xのものの間に挟まってますが（スケートボーダーのものよりも類似していると判定している）、私はそんなに違和感を感じませんでした。

ちなみに画面下部にはCohere Embedの1024次元のベクトルを表示しています。
Titan Embeddingsは1536次元だったのでCohere Embedの方が次元の数は少ないですね。

データセット全量に対するスピアマン順位相関係数を求める

ここまでのところ、Titan Embeddingsと大きな違いはありませんでした。
上述のデータセットの1457件のデータに対して、テストプログラムを作ってCohere Embedのコサイン類似度とデータセットのラベルについて相関係数を求めてみます。

実行

上述のデータセットをカレントディレクトリに保存し、以下を実行します。

spearmanr_cohere.py

from langchain.embeddings import BedrockEmbeddings
from sklearn.metrics.pairwise import cosine_similarity
import json
from scipy.stats import spearmanr

cosine_ary = []
label_ary = []
with open("valid-v1.1.json", encoding="utf-8") as f:
    for line in f:
        line = json.loads(line)

        # 文章1,2のEmbedding
        embeddings = BedrockEmbeddings(model_id="cohere.embed-multilingual-v3")
        vector_1 = embeddings.embed_documents([line["sentence1"]])
        vector_2 = embeddings.embed_documents([line["sentence2"]])

        # コサイン類似度の計算
        similarities = cosine_similarity(vector_1, vector_2).tolist()[0][0]

        # 配列に追加
        cosine_ary.append(similarities)
        label = line["label"]
        label_ary.append(label)
        
# スピアマン順位相関係数
correlation, pvalue = spearmanr(cosine_ary,label_ary)
print(correlation)

実行結果

0.8143659029911076

Cohere Embedのスピアマン順位相関係数は0.814となりました。
Titan Embeddingsは0.729だったので、Cohere Embedの方が（使用したデータセットに対しては）優れてそうです。

他の方も比較結果を見てみても、

Cohere Embed ＞ OpenAI Embeddings ＞ Amazon Titan Embeddings

となっていたので、AWSのBedrock環境でも日本語のEmbeddingsに関してはOpenAIよりも優れている環境が手に入った…と言いたいところですが、Cohere Embedの入力可能トークン数は512のようなので、チャンクの大きさが512トークンで収まるようなユースケースで使用するのが良さそうです。
ちなみに実行速度はTitan Embeddingsと比べると少し速いですね。

日本語でRAGを実装する場合、チャンクの大きさはどれぐらいが適切なんでしょうね。
ドキュメントと検索の仕方によって変わりそうですが。

Cohere Embedのトークン数

1トークン ＞ 1文字のような感じです。

ベッドとソファの置かれた部屋です。(17文字)
"inputTokenCount": 11

カラフルなベッドとソファーが並べて置いてあります。(25文字)
"inputTokenCount": 15

花柄のベッドの上に二匹の猫が座っています。(21文字)
"inputTokenCount": 15

ベッドの布団の上に二匹の猫がいます。(18文字)
"inputTokenCount": 14

階段の手すりでスケートボードをする男性がいます。(24文字)
"inputTokenCount": 16