TF-IDF and Bag of Words: Text Representation & Search › Code Implementation

def tokenize(text):
    """Split text into lowercase words, removing punctuation."""
    import re

    # Convert to lowercase and split on non-word characters
    tokens = re.findall(r"\b\w+\b", text.lower())
    return tokens


## Example documents
documents = [
    "the cat sat on the mat",
    "the dog sat on the log",
    "the cat and dog sat",
]

## Tokenize all documents
tokenized_docs = [tokenize(doc) for doc in documents]
print("Tokenized documents:")
for i, tokens in enumerate(tokenized_docs, 1):
    print(f"Doc {i}: {tokens}")

def build_vocabulary(tokenized_documents):
    """Create a sorted vocabulary from all tokenized documents."""
    # Collect all unique words
    all_words = set()
    for tokens in tokenized_documents:
        all_words.update(tokens)

    # Sort for consistent ordering
    vocabulary = sorted(all_words)

    # Create word-to-index mapping
    word_to_idx = {word: idx for idx, word in enumerate(vocabulary)}

    return vocabulary, word_to_idx


vocabulary, word_to_idx = build_vocabulary(tokenized_docs)
print(f"Vocabulary ({len(vocabulary)} words): {vocabulary}")
print("\nWord to index mapping:")
for word, idx in word_to_idx.items():
    print(f"  {word}: {idx}")

def bag_of_words_vectorize(tokenized_doc, word_to_idx):
    """Convert a tokenized document to a Bag of Words vector."""
    # Initialize vector with zeros
    vector = [0] * len(word_to_idx)

    # Count word occurrences
    for word in tokenized_doc:
        if word in word_to_idx:
            idx = word_to_idx[word]
            vector[idx] += 1

    return vector


## Vectorize all documents
bow_vectors = [
    bag_of_words_vectorize(tokens, word_to_idx) for tokens in tokenized_docs
]

print("Bag of Words vectors:")
for i, vector in enumerate(bow_vectors, 1):
    print(f"Doc {i}: {vector}")

import math


def term_frequency(word, tokenized_doc, use_log=False):
    """Calculate term frequency for a word in a document."""
    count = tokenized_doc.count(word)

    if count == 0:
        return 0.0

    if use_log:
        # Logarithmic scaling
        return 1 + math.log(count)
    else:
        # Normalized by document length
        return count / len(tokenized_doc)


## Calculate TF for "cat" in each document
print("Term Frequency for 'cat':")
for i, tokens in enumerate(tokenized_docs, 1):
    tf_normalized = term_frequency("cat", tokens, use_log=False)
    tf_log = term_frequency("cat", tokens, use_log=True)
    print(f"Doc {i}: normalized={tf_normalized:.3f}, log={tf_log:.3f}")

def inverse_document_frequency(word, tokenized_documents):
    """Calculate IDF for a word across a document collection."""
    # Count documents containing the word
    doc_count = sum(1 for tokens in tokenized_documents if word in tokens)

    if doc_count == 0:
        return 0.0  # Word doesn't appear anywhere

    total_docs = len(tokenized_documents)
    # Standard IDF formula
    idf = math.log(total_docs / doc_count)
    return idf


## Calculate IDF for each word in vocabulary
print("Inverse Document Frequency (IDF):")
idf_scores = {}
for word in vocabulary:
    idf = inverse_document_frequency(word, tokenized_docs)
    idf_scores[word] = idf
    print(
        f"  {word:6s}: {idf:.3f} (appears in {sum(1 for tokens in tokenized_docs if word in tokens)}/{len(tokenized_docs)} docs)"
    )

def tfidf_vectorize(
    tokenized_doc, word_to_idx, tokenized_documents, use_log_tf=False
):
    """Convert a tokenized document to a TF-IDF vector."""
    vector = [0.0] * len(word_to_idx)

    for word, idx in word_to_idx.items():
        # Calculate TF
        tf = term_frequency(word, tokenized_doc, use_log=use_log_tf)

        # Calculate IDF
        idf = inverse_document_frequency(word, tokenized_documents)

        # TF-IDF is the product
        vector[idx] = tf * idf

    return vector


## Compute TF-IDF vectors for all documents
tfidf_vectors = [
    tfidf_vectorize(tokens, word_to_idx, tokenized_docs, use_log_tf=False)
    for tokens in tokenized_docs
]

print("TF-IDF vectors (using normalized TF):")
for i, vector in enumerate(tfidf_vectors, 1):
    print(f"\nDoc {i}:")
    # Show non-zero values for clarity
    non_zero = [
        (vocabulary[j], f"{vector[j]:.3f}")
        for j, val in enumerate(vector)
        if val > 0
    ]
    for word, score in sorted(
        non_zero, key=lambda x: float(x[1]), reverse=True
    ):
        print(f"  {word:6s}: {score}")

def cosine_similarity(vec1, vec2):
    """Calculate cosine similarity between two vectors."""
    import math

    # Dot product
    dot_product = sum(a * b for a, b in zip(vec1, vec2))

    # Magnitudes
    magnitude1 = math.sqrt(sum(a * a for a in vec1))
    magnitude2 = math.sqrt(sum(b * b for b in vec2))

    if magnitude1 == 0 or magnitude2 == 0:
        return 0.0

    return dot_product / (magnitude1 * magnitude2)


## Compare all pairs of documents
print("Cosine similarity between documents (using TF-IDF):")
for i in range(len(tfidf_vectors)):
    for j in range(i + 1, len(tfidf_vectors)):
        similarity = cosine_similarity(tfidf_vectors[i], tfidf_vectors[j])
        print(f"Doc {i + 1} vs Doc {j + 1}: {similarity:.3f}")

from sklearn.feature_extraction.text import TfidfVectorizer

## Initialize vectorizer
vectorizer = TfidfVectorizer(
    lowercase=True,
    token_pattern=r"\b\w+\b",  # Word boundaries
    max_features=1000,  # Limit vocabulary size
    min_df=2,  # Ignore words appearing in < 2 documents
    max_df=0.95,  # Ignore words appearing in > 95% of documents
)

## Fit and transform
tfidf_matrix = vectorizer.fit_transform(documents)

print(f"TF-IDF matrix shape: {tfidf_matrix.shape}")
print(f"Vocabulary size: {len(vectorizer.vocabulary_)}")
print(f"\nSample vocabulary words: {list(vectorizer.vocabulary_.keys())[:10]}")

## Convert to dense array for display (not recommended for large datasets)
dense_matrix = tfidf_matrix.toarray()

print("TF-IDF matrix (documents × words):")
print(f"\nWords: {vectorizer.get_feature_names_out()}")
print(f"\nMatrix:\n{dense_matrix}")