Efficient Text Processing: Python TF IDF Code Explained

We will see a step by step implementation of TF IDF from scratch in Python.

Step 1: Install Required Libraries

pip install nltk numpy

Step 2: Import Libraries and Prepare Data

import nltk
import numpy as np
import math
from collections import Counter

Some sample records used as out corpus 
documents = [
    "the cat in the hat",
    "the quick brown fox",
    "the cat and the hat are friends",
    "the fox jumps over the lazy dog"
]

Tokenizing the documents using nltk

nltk.download('punkt')
tokenized_docs = [nltk.word_tokenize(doc.lower()) for doc in documents]

Step 3: Computing Term Frequency (TF) based on formula

from collections import Counter

def compute_tf(tokenized_docs):
    """
    Compute Term Frequency (TF) for each document in a list of tokenized documents.
    
    Parameters:
    tokenized_docs (list of list of str): A list where each element is a list of tokens (words) for a document.
    
    Returns:
    list of dict: A list of dictionaries, each containing the term frequencies for a corresponding document.
    """
    tf_list = []  # Initialize a list to hold the TF scores for each document
    
    # Iterate through each document's tokens
    for tokens in tokenized_docs:
        # Count the frequency of each token in the document
        count = Counter(tokens)
        
        # Calculate the term frequency for each word
        tf = {word: freq / len(tokens) for word, freq in count.items()}
        
        # Append the TF dictionary to the list
        tf_list.append(tf)
    
    return tf_list  # Return the list of TF scores for all documents

# Example usage of the compute_tf function
tf_scores = compute_tf(tokenized_docs)  # Compute the TF scores for the tokenized documents

# Print the TF scores for each document
print("TF Scores:")
for i, tf in enumerate(tf_scores):
    print(f"Document {i + 1}: {tf}")  # Display the TF scores for each document

Output

TF Scores:
Document 1: {'the': 0.4, 'cat': 0.2, 'in': 0.2, 'hat': 0.2}
Document 2: {'the': 0.25, 'quick': 0.25, 'brown': 0.25, 'fox': 0.25}
Document 3: {'the': 0.2857142857142857, 'cat': 0.14285714285714285, 'and': 0.14285714285714285, 'hat': 0.14285714285714285, 'are': 0.14285714285714285, 'friends': 0.14285714285714285}
Document 4: {'the': 0.2857142857142857, 'fox': 0.14285714285714285, 'jumps': 0.14285714285714285, 'over': 0.14285714285714285, 'lazy': 0.14285714285714285, 'dog': 0.14285714285714285}

Step 4: Computing Inverse Document Frequency based on formula

import math

def compute_idf(tokenized_docs):
    """
    The function computes Inverse Document Frequency (IDF) for each unique word in a list of tokenized documents.
    
    Input parameters of function:
    tokenized_docs (list of list of str): A list where each element is a list of tokens (words) for a document.
    
    Returns:
    dict: A dictionary where keys are words and values are their corresponding IDF scores.
    """
    n = len(tokenized_docs)  # Total number of documents
    idf = {}  # Initialize a dictionary to hold IDF scores
    # Create a set of all unique words across all documents
    all_words = set(word for tokens in tokenized_docs for word in tokens)
    
    # Calculate IDF for each unique word
    for word in all_words:
        # Count how many documents contain the word (document frequency)
        df = sum(1 for tokens in tokenized_docs if word in tokens)
        # Compute the IDF score using the formula: log(n / (1 + df))
        idf[word] = math.log(n / (1 + df))  # Add 1 to avoid division by zero

    return idf  # Return the dictionary of IDF scores

# Example usage of the compute_idf function
idf_scores = compute_idf(tokenized_docs)  # Compute the IDF scores for the tokenized documents

# Print the IDF scores for each word
print("\nIDF Scores:")
print(idf_scores)  # Display the IDF scores

Output

IDF Scores:
{'in': 0.6931471805599453, 'jumps': 0.6931471805599453, 'brown': 0.6931471805599453, 'over': 0.6931471805599453, 'are': 0.6931471805599453, 'cat': 0.28768207245178085,
 'lazy': 0.6931471805599453, 'the': -0.2231435513142097, 'fox': 0.28768207245178085, 'quick': 0.6931471805599453, 'friends': 0.6931471805599453, 'and': 0.6931471805599453, 'hat': 0.28768207245178085, 'dog': 0.6931471805599453}

Step 5: Compute TF-IDF

We can combine the TF and IDF scores to compute TF-IDF:

def compute_tf_idf(tf_scores, idf_scores):
    """
    The function Computes TF-IDF scores for each document using the term frequency (TF) scores
    and the inverse document frequency (IDF) scores.

    Input Parameters of the function:
    tf_scores (list of dict): A list where each dictionary contains the TF scores for a document.
    idf_scores (dict): A dictionary where keys are words and values are their corresponding IDF scores.

    Returns:
    list of dict: A list of dictionaries, each containing the TF-IDF scores for a corresponding document.
    """
    tf_idf_list = []  # Initialize a list to hold the TF-IDF scores for each document
    
    # Iterate through each document's TF scores
    for tf in tf_scores:
        # Compute TF-IDF for each word in the document
        tf_idf = {word: tf_val * idf_scores[word] for word, tf_val in tf.items()}
        # Append the TF-IDF dictionary to the list
        tf_idf_list.append(tf_idf)
    
    return tf_idf_list  # Return the list of TF-IDF scores for all documents

# Example usage of the compute_tf_idf function
tf_idf_scores = compute_tf_idf(tf_scores, idf_scores)  # Compute the TF-IDF scores using the TF and IDF scores

# Print the TF-IDF scores for each document
print("\nTF-IDF Scores:")
for i, tfidf in enumerate(tf_idf_scores):
    print(f"Document {i + 1}: {tfidf}")  # Display the TF-IDF scores for each document

Output

TF-IDF Scores:
Document 1: {'the': -0.08925742052568389, 'cat': 0.05753641449035617, 'in': 0.13862943611198905, 'hat': 0.05753641449035617}
Document 2: {'the': -0.05578588782855243, 'quick': 0.17328679513998632, 'brown': 0.17328679513998632, 'fox': 0.07192051811294521}
Document 3: {'the': -0.06375530037548849, 'cat': 0.04109743892168297, 'and': 0.09902102579427789, 'hat': 0.04109743892168297, 'are': 0.09902102579427789, 'friends': 0.09902102579427789}
Document 4: {'the': -0.06375530037548849, 'fox': 0.04109743892168297, 'jumps': 0.09902102579427789, 'over': 0.09902102579427789, 'lazy': 0.09902102579427789, 'dog': 0.09902102579427789}