TF-IDF Search Engine (python, written by Claude Code)

Claude Code wrote this python project from the task below. It installed and ran on a clean Ubuntu 22.04 machine as written.

Task given to the agent:

TASK: TF-IDF Search Engine

Write a program that builds a TF-IDF (Term Frequency-Inverse Document Frequency) index over a collection of text documents and supports ranked keyword search queries returning the most relevant documents.

FUNCTIONAL REQUIREMENTS:
- Accept a directory of text files as a command-line argument to build the index
- Tokenize documents: split on whitespace and punctuation, convert to lowercase, remove stop words (built-in list of common English stop words like "the", "is", "and", etc.)
- Support optional stemming/lemmatization via --stem flag to group word variants (e.g., "running", "runs", "ran" all map to "run")
- Compute TF-IDF scores for each term in each document using standard formulas: TF = term count / total terms in document, IDF = log(total documents / documents containing term)
- Accept search queries via --query flag and return the top N most relevant documents ranked by cosine similarity between query vector and document vectors (--top flag, default 10)
- Support multi-word queries: compute a query TF-IDF vector and rank documents by similarity
- Support boolean operators in queries via --boolean flag: AND (both terms required), OR (either term), NOT (exclude term)
- Display search results showing: rank, document name, relevance score, and a snippet of the matching text with query terms highlighted
- Save the built index to a file via --save-index flag for reuse without reprocessing
- Load a previously saved index via --load-index flag
- Print index statistics: total documents, total unique terms, average document length, most common terms (top 20)
- Save search results as JSON with --output flag
- If no directory is given, generate a sample corpus of 20 short documents on varied topics (science, sports, technology, cooking, travel), build the index, and demonstrate several search queries with ranked results
- Handle errors: empty documents, binary files in the directory, extremely large documents, and empty queries

Create a complete Python project for a clean Ubuntu 22.04 machine with only Python 3.10+ installed. Include:
- Source code
- requirements.txt with all dependencies (direct and transitive) pinned to exact versions
- README.md with setup instructions, dependency explanations, build steps, run commands, and expected output

[]

# TF-IDF Search Engine (Python - Trial 1)

A search engine that builds a TF-IDF index over a collection of documents and supports ranked keyword search using cosine similarity.

## Dependencies

- **scikit-learn** (1.4.0) - TF-IDF vectorization and cosine similarity computation
- **nltk** (3.8.1) - Tokenization, stopword removal, and Porter stemming

## Setup

```bash
pip install -r requirements.txt
```

NLTK data packages (`punkt`, `stopwords`) are downloaded automatically on first run.

## Usage

### Demo mode (built-in sample documents)
```bash
python search.py --demo
```

### Index a directory of text files
```bash
python search.py --dir /path/to/documents --query "search terms"
```

### Interactive search
```bash
python search.py --demo --interactive
```

### Command-line options
```
--dir DIR          Directory containing .txt files to index
--query, -q QUERY  Search query to execute
--top-k, -k K      Number of top results (default: 5)
--no-stemming      Disable Porter stemming
--interactive, -i  Interactive search mode
--demo             Use built-in sample documents
```

## How It Works

1. Documents are tokenized using NLTK's word tokenizer
2. Stopwords are removed and optional Porter stemming is applied
3. scikit-learn's `TfidfVectorizer` builds the TF-IDF matrix
4. Queries are vectorized and ranked by cosine similarity against indexed documents

scikit-learn==1.4.0
nltk==3.8.1
numpy>=1.24.0

"""
TF-IDF Search Engine

Builds a TF-IDF index over a collection of documents and supports
ranked keyword search using cosine similarity.

Dependencies: scikit-learn (1.4.0), nltk (3.8.1)
"""

import os
import sys
import json
import argparse
from pathlib import Path

import nltk
from nltk.tokenize import word_tokenize
from nltk.corpus import stopwords
from nltk.stem import PorterStemmer
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np


# Ensure NLTK data is available
def ensure_nltk_data():
    """Download required NLTK data packages if not already present."""
    for resource in ["punkt", "punkt_tab", "stopwords"]:
        try:
            nltk.data.find(f"tokenizers/{resource}" if "punkt" in resource else f"corpora/{resource}")
        except LookupError:
            nltk.download(resource, quiet=True)


class TFIDFSearchEngine:
    """A search engine that uses TF-IDF vectorization and cosine similarity."""

    def __init__(self, use_stemming=True, max_features=10000):
        """
        Initialize the search engine.

        Args:
            use_stemming: Whether to apply Porter stemming to tokens.
            max_features: Maximum number of features for the TF-IDF vectorizer.
        """
        ensure_nltk_data()
        self.use_stemming = use_stemming
        self.stemmer = PorterStemmer() if use_stemming else None
        self.stop_words = set(stopwords.words("english"))
        self.documents = []
        self.doc_names = []
        self.vectorizer = TfidfVectorizer(
            max_features=max_features,
            tokenizer=self._tokenize,
            token_pattern=None,
        )
        self.tfidf_matrix = None

    def _tokenize(self, text):
        """
        Tokenize text with optional stemming and stopword removal.

        Args:
            text: The input text to tokenize.

        Returns:
            A list of processed tokens.
        """
        tokens = word_tokenize(text.lower())
        tokens = [t for t in tokens if t.isalnum() and t not in self.stop_words]
        if self.stemmer:
            tokens = [self.stemmer.stem(t) for t in tokens]
        return tokens

    def add_document(self, name, content):
        """
        Add a document to the collection.

        Args:
            name: The document identifier or filename.
            content: The text content of the document.
        """
        self.doc_names.append(name)
        self.documents.append(content)

    def load_documents_from_directory(self, directory):
        """
        Load all .txt files from a directory.

        Args:
            directory: Path to the directory containing text files.
        """
        dir_path = Path(directory)
        if not dir_path.is_dir():
            raise FileNotFoundError(f"Directory not found: {directory}")

        for filepath in sorted(dir_path.glob("*.txt")):
            content = filepath.read_text(encoding="utf-8", errors="replace")
            self.add_document(filepath.name, content)

        print(f"Loaded {len(self.documents)} document(s) from '{directory}'.")

    def build_index(self):
        """Build the TF-IDF index from the loaded documents."""
        if not self.documents:
            raise ValueError("No documents to index. Add documents first.")

        self.tfidf_matrix = self.vectorizer.fit_transform(self.documents)
        print(f"Index built: {self.tfidf_matrix.shape[0]} documents, "
              f"{self.tfidf_matrix.shape[1]} terms.")

    def search(self, query, top_k=10):
        """
        Search the index with a keyword query.

        Args:
            query: The search query string.
            top_k: Number of top results to return.

        Returns:
            A list of (document_name, score) tuples sorted by relevance.
        """
        if self.tfidf_matrix is None:
            raise ValueError("Index not built. Call build_index() first.")

        query_vector = self.vectorizer.transform([query])
        similarities = cosine_similarity(query_vector, self.tfidf_matrix).flatten()

        # Get top-k indices sorted by descending score
        top_indices = similarities.argsort()[::-1][:top_k]

        results = []
        for idx in top_indices:
            score = similarities[idx]
            if score > 0:
                results.append((self.doc_names[idx], float(score)))

        return results

    def get_vocabulary_size(self):
        """Return the number of terms in the vocabulary."""
        if self.vectorizer.vocabulary_:
            return len(self.vectorizer.vocabulary_)
        return 0

    def get_document_count(self):
        """Return the number of indexed documents."""
        return len(self.documents)

    def save_index(self, filepath):
        """
        Save the index metadata to a JSON file.

        Args:
            filepath: Path to the output file.
        """
        metadata = {
            "document_count": len(self.documents),
            "vocabulary_size": self.get_vocabulary_size(),
            "document_names": self.doc_names,
            "use_stemming": self.use_stemming,
        }
        with open(filepath, "w", encoding="utf-8") as f:
            json.dump(metadata, f, indent=2)
        print(f"Index metadata saved to '{filepath}'.")


def create_sample_documents():
    """Create sample documents for demonstration purposes."""
    return {
        "python_intro.txt": (
            "Python is a high-level programming language known for its simplicity "
            "and readability. It supports multiple programming paradigms including "
            "object-oriented, procedural, and functional programming. Python has a "
            "large standard library and an active community."
        ),
        "machine_learning.txt": (
            "Machine learning is a subset of artificial intelligence that enables "
            "systems to learn from data. Supervised learning uses labeled data to "
            "train models, while unsupervised learning finds patterns in unlabeled "
            "data. Deep learning uses neural networks with many layers."
        ),
        "web_development.txt": (
            "Web development encompasses building websites and web applications. "
            "Frontend development focuses on the user interface using HTML, CSS, "
            "and JavaScript. Backend development handles server-side logic, "
            "databases, and APIs."
        ),
        "data_science.txt": (
            "Data science combines statistics, mathematics, and computer science "
            "to extract insights from data. Common tools include Python, R, and "
            "SQL. Data visualization helps communicate findings effectively."
        ),
        "algorithms.txt": (
            "Algorithms are step-by-step procedures for solving problems. Sorting "
            "algorithms like quicksort and mergesort organize data efficiently. "
            "Search algorithms such as binary search find elements in collections. "
            "Graph algorithms traverse networks and find shortest paths."
        ),
    }


def main():
    """Main entry point for the TF-IDF search engine CLI."""
    parser = argparse.ArgumentParser(
        description="TF-IDF Search Engine - Build an index and search documents."
    )
    parser.add_argument(
        "--dir",
        type=str,
        default=None,
        help="Directory containing .txt files to index.",
    )
    parser.add_argument(
        "--query", "-q",
        type=str,
        default=None,
        help="Search query to execute.",
    )
    parser.add_argument(
        "--top-k", "-k",
        type=int,
        default=5,
        help="Number of top results to return (default: 5).",
    )
    parser.add_argument(
        "--no-stemming",
        action="store_true",
        help="Disable Porter stemming.",
    )
    parser.add_argument(
        "--interactive", "-i",
        action="store_true",
        help="Run in interactive mode.",
    )
    parser.add_argument(
        "--demo",
        action="store_true",
        help="Run with built-in sample documents.",
    )

    args = parser.parse_args()

    engine = TFIDFSearchEngine(use_stemming=not args.no_stemming)

    # Load documents
    if args.demo or (args.dir is None and args.query is None and not args.interactive):
        print("Running with sample documents...\n")
        samples = create_sample_documents()
        for name, content in samples.items():
            engine.add_document(name, content)
    elif args.dir:
        engine.load_documents_from_directory(args.dir)
    else:
        print("No document source specified. Use --dir, --demo, or provide a query with --demo.")
        sys.exit(1)

    # Build the index
    engine.build_index()
    print(f"Vocabulary size: {engine.get_vocabulary_size()}")
    print()

    # Execute query or enter interactive mode
    if args.interactive:
        print("Interactive search mode. Type 'quit' or 'exit' to stop.\n")
        while True:
            try:
                query = input("Query> ").strip()
            except (EOFError, KeyboardInterrupt):
                print("\nExiting.")
                break

            if query.lower() in ("quit", "exit", "q"):
                print("Exiting.")
                break

            if not query:
                continue

            results = engine.search(query, top_k=args.top_k)
            if results:
                print(f"\nTop {len(results)} result(s) for '{query}':")
                for rank, (name, score) in enumerate(results, 1):
                    print(f"  {rank}. {name} (score: {score:.4f})")
            else:
                print(f"No results found for '{query}'.")
            print()
    elif args.query:
        results = engine.search(args.query, top_k=args.top_k)
        if results:
            print(f"Top {len(results)} result(s) for '{args.query}':")
            for rank, (name, score) in enumerate(results, 1):
                print(f"  {rank}. {name} (score: {score:.4f})")
        else:
            print(f"No results found for '{args.query}'.")
    else:
        # Default: demo with a few sample queries
        sample_queries = [
            "machine learning neural networks",
            "python programming language",
            "web frontend JavaScript",
            "sorting search algorithms",
        ]
        for query in sample_queries:
            results = engine.search(query, top_k=3)
            print(f"Query: '{query}'")
            if results:
                for rank, (name, score) in enumerate(results, 1):
                    print(f"  {rank}. {name} (score: {score:.4f})")
            else:
                print("  No results found.")
            print()


if __name__ == "__main__":
    main()