sparse_dot_topnを使った高速なcos類似度計算

pip install sparse_dot_topn

import numpy as np
from scipy.sparse import csr_matrix
from scipy.sparse import rand
from sparse_dot_topn import awesome_cossim_topn

N = 10
a = rand(100, 1000000, density=0.005, format='csr')
b = rand(1000000, 200, density=0.005, format='csr')

# Default precision type is np.float64, but you can down cast to have a small memory footprint and faster execution
# Remark : These are the only 2 types supported now, since we assume that float16 will be difficult to implement and will be slower, because C doesn't support a 16-bit float type on most PCs
a = a.astype(np.float32)
b = b.astype(np.float32)

# Use standard implementation
c = awesome_cossim_topn(a, b, N, 0.01)

# Use parallel implementation with 4 threads
d = awesome_cossim_topn(a, b, N, 0.01, use_threads=True, n_jobs=4)

# Use standard implementation with 4 threads and with the computation of best_ntop: the value of ntop needed to capture all results above lower_bound
d, best_ntop = awesome_cossim_topn(a, b, N, 0.01, use_threads=True, n_jobs=4, return_best_ntop=True)

from tqdm.auto import tqdm
from datasets import load_dataset
from janome.tokenizer import Tokenizer

dataset = load_dataset("shunk031/JGLUE", 'JSTS', split='train')

tokenizer = Tokenizer()

queries, corpus = [], []
for d in tqdm(dataset):
    if d['label'] < 4:
        continue
    s1 = " ".join([token.surface for token in tokenizer.tokenize(d['sentence1'])])
    s2 = " ".join([token.surface for token in tokenizer.tokenize(d['sentence2'])])

    queries.append(s1)
    corpus.append(s2)

import time
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
from sparse_dot_topn import awesome_cossim_topn


def calc_cos_for_retrieve(corpus, queries, mode):
    vectorizer = TfidfVectorizer()
    X = vectorizer.fit_transform(corpus)
    
    t0 = time.time()
    p_at_n = []
    for i, query in enumerate(queries):
        qx = vectorizer.transform([query])

        if mode == "sklearn":
            sim = cosine_similarity(X, qx, dense_output=False)
        elif mode == "sparse_dot_topn":
            sim = awesome_cossim_topn(X, qx.T, 10)
        else:
            raise ValueError(f"Unknown mode `{mode}`")
        sim = sim.toarray()[:, 0]
        
        relevant_documents = np.argsort(-sim)
        idx = np.where(relevant_documents == i)
        idx = idx[0][0]
        
        p_at_n.append(idx < 10)  # P@10
    t = time.time() - t0
    
    return np.mean(p_at_n), t

p_at_n, _ = calc_cos_for_retrieve(corpus, queries, "sklearn")
p_at_n  # 0.775570272259014

p_at_n, _ = calc_cos_for_retrieve(corpus, queries, "sparse_dot_topn")
p_at_n  # 0.775570272259014