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

Claude Code wrote this cpp 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 C++ project for a clean Ubuntu 22.04 machine with only G++ 12+ and CMake 3.22+ installed. Include:
- Source code
- CMakeLists.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

[]

cmake_minimum_required(VERSION 3.22)
project(tfidf_search VERSION 1.0.0 LANGUAGES CXX)

set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_CXX_EXTENSIONS OFF)

include(FetchContent)

FetchContent_Declare(
    eigen
    GIT_REPOSITORY https://gitlab.com/libeigen/eigen.git
    GIT_TAG        3.4.0
    GIT_SHALLOW    TRUE
)

FetchContent_Declare(
    nlohmann_json
    GIT_REPOSITORY https://github.com/nlohmann/json.git
    GIT_TAG        v3.11.3
    GIT_SHALLOW    TRUE
)

set(EIGEN_BUILD_DOC OFF CACHE BOOL "" FORCE)
set(EIGEN_BUILD_TESTING OFF CACHE BOOL "" FORCE)
set(JSON_BuildTests OFF CACHE BOOL "" FORCE)
set(JSON_Install OFF CACHE BOOL "" FORCE)

FetchContent_MakeAvailable(eigen nlohmann_json)

add_executable(tfidf_search main.cpp)

target_link_libraries(tfidf_search PRIVATE
    Eigen3::Eigen
    nlohmann_json::nlohmann_json
)

if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang")
    target_compile_options(tfidf_search PRIVATE -Wall -Wextra -Wpedantic)
endif()

/**
 * TF-IDF Search Engine (C++ - Trial 1)
 *
 * Builds a TF-IDF index over documents and supports ranked keyword search
 * using cosine similarity with Eigen for linear algebra.
 *
 * Dependencies:
 *   - Eigen (3.4.0) - Linear algebra for TF-IDF matrix and cosine similarity
 *   - nlohmann/json (3.11.3) - JSON metadata output
 */

#include <Eigen/Dense>
#include <Eigen/Sparse>
#include <nlohmann/json.hpp>

#include <algorithm>
#include <cmath>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <map>
#include <regex>
#include <set>
#include <sstream>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

namespace fs = std::filesystem;
using json = nlohmann::json;

static const std::unordered_set<std::string> STOP_WORDS = {
    "a","an","the","is","are","was","were","be","been","being","have","has","had",
    "do","does","did","will","would","could","should","may","might","shall","can",
    "to","of","in","for","on","with","at","by","from","as","into","through","during",
    "before","after","above","below","between","out","off","over","under","again",
    "then","once","here","there","when","where","why","how","all","both","each",
    "few","more","most","other","some","such","no","nor","not","only","own","same",
    "so","than","too","very","just","because","but","and","or","if","while","about",
    "up","it","its","this","that","these","those","i","me","my","we","our","you",
    "your","he","him","his","she","her","they","them","their","what","which","who",
};

static std::string to_lower(const std::string& s) {
    std::string out = s;
    std::transform(out.begin(), out.end(), out.begin(),
                   [](unsigned char c) { return std::tolower(c); });
    return out;
}

static std::vector<std::string> tokenize(const std::string& text) {
    std::vector<std::string> tokens;
    std::regex word_re("[a-zA-Z0-9]+");
    auto begin = std::sregex_iterator(text.begin(), text.end(), word_re);
    auto end = std::sregex_iterator();
    for (auto it = begin; it != end; ++it) {
        std::string tok = to_lower(it->str());
        if (tok.size() > 1 && STOP_WORDS.find(tok) == STOP_WORDS.end()) {
            tokens.push_back(tok);
        }
    }
    return tokens;
}

static std::string read_file(const std::string& path) {
    std::ifstream ifs(path, std::ios::binary);
    if (!ifs) throw std::runtime_error("Cannot open: " + path);
    std::ostringstream ss;
    ss << ifs.rdbuf();
    return ss.str();
}

class TFIDFSearchEngine {
public:
    void add_document(const std::string& name, const std::string& content) {
        doc_names_.push_back(name);
        documents_.push_back(content);
    }

    void load_from_directory(const std::string& dir_path) {
        std::vector<fs::path> files;
        for (const auto& entry : fs::directory_iterator(dir_path)) {
            if (entry.path().extension() == ".txt") files.push_back(entry.path());
        }
        std::sort(files.begin(), files.end());
        for (const auto& f : files) {
            add_document(f.filename().string(), read_file(f.string()));
        }
        std::cout << "Loaded " << files.size() << " document(s) from '" << dir_path << "'.\n";
    }

    void build_index() {
        if (documents_.empty()) throw std::runtime_error("No documents to index.");

        // Build vocabulary
        std::map<std::string, int> vocab_map;
        std::vector<std::vector<std::string>> doc_tokens;
        for (const auto& doc : documents_) {
            auto tokens = tokenize(doc);
            for (const auto& t : tokens) {
                if (vocab_map.find(t) == vocab_map.end()) {
                    int idx = static_cast<int>(vocab_map.size());
                    vocab_map[t] = idx;
                }
            }
            doc_tokens.push_back(std::move(tokens));
        }
        vocab_size_ = static_cast<int>(vocab_map.size());
        int num_docs = static_cast<int>(documents_.size());

        // Compute document frequency
        Eigen::VectorXd df = Eigen::VectorXd::Zero(vocab_size_);
        for (const auto& tokens : doc_tokens) {
            std::unordered_set<std::string> unique_terms(tokens.begin(), tokens.end());
            for (const auto& t : unique_terms) {
                df(vocab_map[t]) += 1.0;
            }
        }

        // Compute IDF: log(N / df) + 1
        Eigen::VectorXd idf(vocab_size_);
        for (int i = 0; i < vocab_size_; ++i) {
            idf(i) = std::log(static_cast<double>(num_docs) / (df(i) + 1.0)) + 1.0;
        }

        // Build TF-IDF matrix (docs x terms)
        tfidf_matrix_ = Eigen::MatrixXd::Zero(num_docs, vocab_size_);
        for (int d = 0; d < num_docs; ++d) {
            std::unordered_map<std::string, int> tf;
            for (const auto& t : doc_tokens[d]) tf[t]++;
            for (const auto& [term, count] : tf) {
                int idx = vocab_map[term];
                tfidf_matrix_(d, idx) = static_cast<double>(count) * idf(idx);
            }
            // Normalize row
            double norm = tfidf_matrix_.row(d).norm();
            if (norm > 0) tfidf_matrix_.row(d) /= norm;
        }

        vocab_map_ = std::move(vocab_map);
        idf_ = std::move(idf);
        index_built_ = true;

        std::cout << "Index built: " << num_docs << " documents, "
                  << vocab_size_ << " terms.\n";
    }

    struct SearchResult {
        std::string name;
        double score;
    };

    std::vector<SearchResult> search(const std::string& query, int top_k = 5) const {
        if (!index_built_) throw std::runtime_error("Index not built.");

        auto tokens = tokenize(query);
        Eigen::VectorXd query_vec = Eigen::VectorXd::Zero(vocab_size_);
        for (const auto& t : tokens) {
            auto it = vocab_map_.find(t);
            if (it != vocab_map_.end()) {
                query_vec(it->second) += 1.0 * idf_(it->second);
            }
        }

        double qnorm = query_vec.norm();
        if (qnorm > 0) query_vec /= qnorm;

        // Cosine similarity: matrix * query_vec
        Eigen::VectorXd scores = tfidf_matrix_ * query_vec;

        // Get top-k
        std::vector<std::pair<int, double>> indexed_scores;
        for (int i = 0; i < scores.size(); ++i) {
            if (scores(i) > 0) indexed_scores.push_back({i, scores(i)});
        }
        std::sort(indexed_scores.begin(), indexed_scores.end(),
                  [](const auto& a, const auto& b) { return a.second > b.second; });

        std::vector<SearchResult> results;
        for (int i = 0; i < std::min(top_k, static_cast<int>(indexed_scores.size())); ++i) {
            results.push_back({doc_names_[indexed_scores[i].first],
                              indexed_scores[i].second});
        }
        return results;
    }

    int get_document_count() const { return static_cast<int>(documents_.size()); }
    int get_vocabulary_size() const { return vocab_size_; }

    void save_metadata(const std::string& path) const {
        json meta;
        meta["document_count"] = documents_.size();
        meta["vocabulary_size"] = vocab_size_;
        meta["document_names"] = doc_names_;
        std::ofstream ofs(path);
        ofs << meta.dump(2);
        std::cout << "Index metadata saved to '" << path << "'.\n";
    }

    static std::vector<std::pair<std::string, std::string>> create_sample_documents() {
        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."},
            {"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."},
            {"web_development.txt", "Web development encompasses building websites and web applications. Frontend development focuses on the user interface using HTML, CSS, and JavaScript."},
            {"data_science.txt", "Data science combines statistics, mathematics, and computer science to extract insights from data. Common tools include Python, R, and SQL."},
            {"algorithms.txt", "Algorithms are step-by-step procedures for solving problems. Sorting algorithms like quicksort and mergesort organize data efficiently."},
        };
    }

private:
    std::vector<std::string> doc_names_;
    std::vector<std::string> documents_;
    std::map<std::string, int> vocab_map_;
    Eigen::MatrixXd tfidf_matrix_;
    Eigen::VectorXd idf_;
    int vocab_size_ = 0;
    bool index_built_ = false;
};

int main(int argc, char* argv[]) {
    std::string dir_path, query_str;
    int top_k = 5;
    bool demo = false, interactive = false;

    for (int i = 1; i < argc; ++i) {
        std::string arg = argv[i];
        if (arg == "--dir" && i + 1 < argc) dir_path = argv[++i];
        else if ((arg == "--query" || arg == "-q") && i + 1 < argc) query_str = argv[++i];
        else if ((arg == "--top-k" || arg == "-k") && i + 1 < argc) top_k = std::stoi(argv[++i]);
        else if (arg == "--demo") demo = true;
        else if (arg == "--interactive" || arg == "-i") interactive = true;
        else if (arg == "--help" || arg == "-h") {
            std::cout << "TF-IDF Search Engine\n\n"
                      << "Options:\n"
                      << "  --dir DIR      Directory of .txt files\n"
                      << "  --query, -q    Search query\n"
                      << "  --top-k, -k    Number of results (default: 5)\n"
                      << "  --demo         Use sample documents\n"
                      << "  --interactive  Interactive search mode\n";
            return 0;
        }
    }

    try {
        TFIDFSearchEngine engine;

        if (!dir_path.empty()) {
            engine.load_from_directory(dir_path);
        } else {
            if (!demo && query_str.empty() && !interactive) demo = true;
            std::cout << "Running with sample documents...\n\n";
            for (auto& [name, content] : TFIDFSearchEngine::create_sample_documents())
                engine.add_document(name, content);
        }

        engine.build_index();
        std::cout << "Vocabulary size: " << engine.get_vocabulary_size() << "\n\n";

        if (interactive) {
            std::cout << "Interactive search mode. Type 'quit' to stop.\n\n";
            std::string line;
            while (true) {
                std::cout << "Query> ";
                if (!std::getline(std::cin, line) || line == "quit" || line == "exit") break;
                if (line.empty()) continue;
                auto results = engine.search(line, top_k);
                if (!results.empty()) {
                    std::cout << "\nTop " << results.size() << " result(s) for '" << line << "':\n";
                    for (size_t i = 0; i < results.size(); ++i)
                        std::cout << "  " << i + 1 << ". " << results[i].name
                                  << " (score: " << results[i].score << ")\n";
                } else {
                    std::cout << "No results found.\n";
                }
                std::cout << "\n";
            }
        } else if (!query_str.empty()) {
            auto results = engine.search(query_str, top_k);
            if (!results.empty()) {
                std::cout << "Top " << results.size() << " result(s) for '" << query_str << "':\n";
                for (size_t i = 0; i < results.size(); ++i)
                    std::cout << "  " << i + 1 << ". " << results[i].name
                              << " (score: " << results[i].score << ")\n";
            } else {
                std::cout << "No results found.\n";
            }
        } else {
            const char* queries[] = {"machine learning neural networks",
                                     "python programming language",
                                     "web frontend JavaScript",
                                     "sorting search algorithms"};
            for (const auto& q : queries) {
                auto results = engine.search(q, 3);
                std::cout << "Query: '" << q << "'\n";
                for (size_t i = 0; i < results.size(); ++i)
                    std::cout << "  " << i + 1 << ". " << results[i].name
                              << " (score: " << results[i].score << ")\n";
                std::cout << "\n";
            }
        }
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << "\n";
        return 1;
    }
    return 0;
}

# TF-IDF Search Engine (C++ - Trial 1)

A TF-IDF search engine using Eigen for linear algebra operations and nlohmann/json for metadata output.

## Dependencies

- **Eigen** (3.4.0) - Linear algebra library for TF-IDF matrix operations and cosine similarity computation
- **nlohmann/json** (3.11.3) - Header-only JSON library for metadata serialization

## Build

```bash
cmake -B build -DCMAKE_BUILD_TYPE=Release
cmake --build build --parallel
```

## Usage

```bash
# Demo with sample documents
./build/tfidf_search --demo

# Index a directory and search
./build/tfidf_search --dir /path/to/docs --query "search terms"

# Interactive mode
./build/tfidf_search --demo --interactive
```

## How It Works

1. Documents are tokenized with stopword removal
2. TF-IDF matrix is built using Eigen dense matrices
3. Queries are vectorized and ranked by cosine similarity
4. Results are sorted by descending relevance score