TF-IDF

Code

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

docs = [
    "the cat sat on the mat",
    "the dog sat on the log",
    "cats and dogs are great pets",
]

# Basic TF-IDF
vec = TfidfVectorizer()
X = vec.fit_transform(docs)
print(X.shape, list(vec.vocabulary_.keys())[:5])

# With n-grams and stopword removal
vec2 = TfidfVectorizer(ngram_range=(1, 2), stop_words="english",
                       max_features=1000, sublinear_tf=True)
X2 = vec2.fit_transform(docs)

# Convert back to terms
feature_names = vec2.get_feature_names_out()
top = feature_names[X2[0].toarray()[0].argsort()[::-1][:3]]
print("top terms:", top)

# Similarity between documents
sim = cosine_similarity(X)
print(sim)

Explanation

TfidfVectorizer converts a corpus into a term-document matrix weighted by term frequency and inverse document frequency. ngram_range captures phrases like 'new york' in addition to single words, and stop_words removes common noise. cosine_similarity over the resulting vectors measures document-to-document similarity.

More NLP Snippets