Python sklearn TfidfVectorizer: Vectorize documents ahead of query for semantic search

user_5307995 0
Viewed 390

I want to run semantic search using TF-IDF.

This code works, but it is really slow when used on a large corpus of documents:

search_terms = "my query"
documents = ["my","list","of","docs"]

vectorizer = TfidfVectorizer()
doc_vectors = vectorizer.fit_transform([search_terms] + documents)
cosine_similarities = linear_kernel(doc_vectors[0:1], doc_vectors).flatten()
document_scores = [item.item() for item in cosine_similarities[1:]]

It seems quite inefficient:

Every new search query triggers a re-vectorizing of the entire corpus.

I am wondering how I can do the bulk work of vectorizing my corpus ahead of time, saving the result in an "index file". So that, when I run a query, the only thing left to do is to vectorize the few words from the query, and then to calculate similarity.

I tried vectorizing query and documents separately:

vec_docs = vectorizer.fit_transform(documents)
vec_query = vectorizer.fit_transform([search_terms])

cosine_similarities = linear_kernel(vec_query, vec_docs).flatten()

But it gives me this error:

ValueError: Incompatible dimension for X and Y matrices: X.shape[1] == 3 while Y.shape[1] == 260541

How can I run the corpus vectorization ahead of time without knowing what the query will be?

My main goal is to get blazing fast results even with a large corpus of documents (say, a few GB worth of text), even on a low-powered server, by doing the bulk of the data-crunching ahead of time.

tfidfvectorizer vectorization python
2 Answers

TF/IDF vectors are high-dimensional and sparse. The basic data structure that supports that is an inverted index. You can either implement it yourself or use a standard index (e.g., Lucene).

Nevertheless, if you would like to experiment with modern deep-neural-based vector representations, check out the following semantic search demo. It uses a similarity search service that can handle billions of vectors.

(Note, I am a co-author of this demo.)

user_159340980

You almost have it right.

In this instance, you can get away with fitting (and transforming) your documents and only transforming your search terms. Here is your code, modified accordingly and using the twenty_newsgroups documents (11k) in its place. You can run it as a script and interactively verify you get fast results:

import numpy as np
from sklearn.datasets import fetch_20newsgroups
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import linear_kernel

news = fetch_20newsgroups()

search_terms = "my query"
# documents = ["my", "list", "of", "docs"]
documents = news.data

vectorizer = TfidfVectorizer()


# fit_transform does two things: fits the vectorizer and transforms documents
doc_vectors = vectorizer.fit_transform(documents)

# the vectorizer is already fit; just transform search_terms via vectorizer
search_term_vector = vectorizer.transform([search_terms])
cosine_similarities = linear_kernel(doc_vectors, search_term_vector).flatten()


if __name__ == "__main__":

    while True:
        query_str = input("\n\n\n\nquery string (return to quit): ")

        if not query_str:
            print("bye!")
            break
        search_term_vector = vectorizer.transform([query_str])

        cosine_similarities = linear_kernel(doc_vectors, search_term_vector).flatten()
        best_idx = np.argmax(cosine_similarities)
        best_score = cosine_similarities[best_idx]
        best_doc = documents[best_idx]
        if best_score < 0.1:
            print("no good matches")
        else:
            max_doc = documents[np.argmax(cosine_similarities)]
            print(
                f"Best match ({round(best_score, 4)}):\n\n", best_doc[0:200] + "...",
            )

Example output:

query string (return to quit): protocol
Best match 0.239 (0.014 sec):

 From: ethan@cs.columbia.edu (Ethan Solomita)
Subject: Re: X protocol packet type
Article-I.D.: cs.C52I2q.IFJ
Organization: Columbia University Department of Computer Science
Lines: 7

In article <9309...

Note: this algorithm find the best match(es) at best in O(n_documents) time, compared to Lucene (powers Elasticsearch) that uses skip lists that can search in O(log(n_documents)). Production search engines also have quiet a bit of tuning to optimize performance. The above could be useful with some tweaking but isn't going to topple Google tomorrow :)

user_50329410
Related