日本語文書の意味的類似度計算が遅い・精度が低い問題をSentence BERTで解決する方法

# TF-IDFによる類似度計算の例
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

# 以下のような文書では類似度が正しく計算されない
doc1 = "機械学習の精度を向上させる方法"
doc2 = "AIモデルの性能を改善する手法"
# 結果: 低い類似度（単語が異なるため）

# BERTの直接利用も計算コストが高い
# 毎回エンコードが必要で、大量データに不向き

pip install sentence-transformers
pip install transformers fugashi ipadic

from transformers import BertJapaneseTokenizer, BertModel
import torch

class SentenceBertJapanese:
    def __init__(self, model_name_or_path, device=None):
        self.tokenizer = BertJapaneseTokenizer.from_pretrained(model_name_or_path)
        self.model = BertModel.from_pretrained(model_name_or_path)
        self.model.eval()

        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        self.device = torch.device(device)
        self.model.to(device)

    def _mean_pooling(self, model_output, attention_mask):
        token_embeddings = model_output[0]
        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)

    @torch.no_grad()
    def encode(self, sentences, batch_size=8):
        all_embeddings = []
        iterator = range(0, len(sentences), batch_size)
        for batch_idx in iterator:
            batch = sentences[batch_idx:batch_idx + batch_size]
            
            encoded_input = self.tokenizer.batch_encode_plus(batch, padding="longest", 
                                           truncation=True, return_tensors="pt").to(self.device)
            model_output = self.model(**encoded_input)
            sentence_embeddings = self._mean_pooling(model_output, encoded_input["attention_mask"]).to('cpu')
            
            all_embeddings.extend(sentence_embeddings)
        
        return torch.stack(all_embeddings)

# 使用例
MODEL_NAME = "sonoisa/sentence-bert-base-ja-mean-tokens-v2"
model = SentenceBertJapanese(MODEL_NAME)

from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

# 文書例
documents = [
    "機械学習の精度を向上させる方法",
    "AIモデルの性能を改善する手法",
    "深層学習の最適化技術",
    "今日の天気は晴れです"
]

# 分散表現の取得
embeddings = model.encode(documents)

# 類似度行列の計算
similarity_matrix = cosine_similarity(embeddings)

print("類似度行列:")
print(similarity_matrix)
# 結果: doc1とdoc2は高い類似度を示す

class DocumentSearchEngine:
    def __init__(self, model):
        self.model = model
        self.document_embeddings = None
        self.documents = None
    
    def index_documents(self, documents):
        """文書の事前エンコード（一度だけ実行）"""
        self.documents = documents
        self.document_embeddings = self.model.encode(documents)
    
    def search(self, query, top_k=5):
        """クエリに類似した文書を高速検索"""
        query_embedding = self.model.encode([query])
        similarities = cosine_similarity(query_embedding, self.document_embeddings)[0]
        
        # 上位k件を取得
        top_indices = np.argsort(similarities)[::-1][:top_k]
        
        results = []
        for idx in top_indices:
            results.append({
                'document': self.documents[idx],
                'similarity': similarities[idx]
            })
        
        return results

# 使用例
search_engine = DocumentSearchEngine(model)
search_engine.index_documents(large_document_corpus)

# 高速検索（ミリ秒オーダー）
results = search_engine.search("機械学習の性能向上", top_k=3)

def recommend_documents(user_document, document_corpus, top_k=5):
    search_engine = DocumentSearchEngine(model)
    search_engine.index_documents(document_corpus)
    return search_engine.search(user_document, top_k)

def detect_duplicates(documents, threshold=0.8):
    embeddings = model.encode(documents)
    similarity_matrix = cosine_similarity(embeddings)
    
    duplicates = []
    for i in range(len(documents)):
        for j in range(i+1, len(documents)):
            if similarity_matrix[i][j] > threshold:
                duplicates.append((i, j, similarity_matrix[i][j]))
    
    return duplicates

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

# 独自データでの追加学習
train_examples = []
for index, row in custom_df.iterrows():
    label = row['category_id']
    train_examples.append(InputExample(texts=[row.text], label=float(label)))

# 学習の実行
train_dataloader = DataLoader(train_examples, batch_size=16, shuffle=True)
train_loss = losses.BatchAllTripletLoss(model)
model.fit(train_objectives=[(train_dataloader, train_loss)], epochs=30)