チャンキング戦略検証用UI 2つの文章の類似度を比較する

概要

　RAGではベクトル検索のチャンキング戦略が重要になる。
　一番手軽なチャンクサイズとオーバーラップによる分割は、文章を余すことなくベクトル化可能だが、文脈が分断された文を含んだ文章でベクトルするという欠点がある。
　オーバーラップさせる目的からも、文脈が分断されて意味が通じない箇所は無視して、分断されていない意味が通じる箇所だけが有効であるほうが、ベクトル検索やその検索でヒットした箇所を入力に使うRAGでは嬉しい挙動だろう。
　今回は、途切れた文章を含んでしまうことがどの程度影響あるのかを確認する目的で、intfloat/multilingual-e5-largeで2つの文章をベクトル化して類似度を確認できるシステムを構築した。

今回のコードでできること

2つの文章を入力するとコサイン類似度を取得できる

intfloat/multilingual-e5-large使用時の注意

• 公式情報(特にFAQ、Limitations)は確認すべきである
• 検索時は前に"query: "を付ける
• DBに格納する文章をベクトル化する時は前に"passage: "を付ける
• 先頭から最大512tokenまでしか使われない。513 token 以降は完全に捨てられる

簡易実験

以下の3文を加工して実験する。

途切れた文章が前後にある場合の影響

###　実験1

テキスト①は、文脈が分断された文を前後に含んだ文章。
テキスト②は、文脈が分断された文を除去した文章。
結果は0.9319。
完全に一致する場合は1になるので、1→0.9319の減少が途切れた文を前後に含む影響といえる。

実験2

テキスト①は、文脈が分断された文を前後に含んだ文章。
テキスト②は、テキスト①の文脈が分断される前の文章。
結果は0.9565。
実験１の時より類似度が高いので、文脈が分断された文でも無視されずベクトルに影響を与えていることが確認され、嬉しくない結果となった。

ソースコード（Pythonファイル）

import gradio as gr
import torch
import torch.nn.functional as F
from torch import Tensor
from transformers import AutoTokenizer, AutoModel

# =========================
# Pooling
# =========================
def average_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
    last_hidden = last_hidden_states.masked_fill(
        ~attention_mask[..., None].bool(), 0.0
    )
    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]

# =========================
# Model Load
# =========================
MODEL_NAME = "intfloat/multilingual-e5-large"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModel.from_pretrained(MODEL_NAME)
model.eval()

# =========================
# Similarity Function
# =========================
def calc_similarity(text1, text2):
    if not text1 or not text2:
        return 0.0

    inputs = tokenizer(
        [text1, text2],
        max_length=512,
        padding=True,
        truncation=True,
        return_tensors="pt"
    )

    with torch.no_grad():
        outputs = model(**inputs)
        embeddings = average_pool(
            outputs.last_hidden_state,
            inputs["attention_mask"]
        )
        embeddings = F.normalize(embeddings, p=2, dim=1)

    # cosine similarity
    score = torch.matmul(embeddings[0], embeddings[1]).item()
    return score

# =========================
# Gradio UI
# =========================
with gr.Blocks(title="Text Similarity (e5-large / transformers)") as demo:
    gr.Markdown("## 🔍 テキスト類似度計測（multilingual-e5-large）")

    with gr.Row():
        t1 = gr.Textbox(label="テキスト①", lines=6)
        t2 = gr.Textbox(label="テキスト②", lines=6)

    btn = gr.Button("類似度を計算")
    out = gr.Number(label="Cosine Similarity", precision=4)

    btn.click(calc_similarity, [t1, t2], out)

demo.launch(debug=True, share=False, inbrowser=False)