Vector Search with Text Embeddings

A simple learning project demonstrating how to use text embeddings for semantic search.

Overview

This project shows how to:

• Compute text embeddings using pre-trained models
• Store embeddings for efficient retrieval
• Perform vector similarity search using cosine similarity

Setup

1. Create and activate a virtual environment:

python3 -m venv venv
source venv/bin/activate

2. Install dependencies:

pip install -r requirements.txt

Usage

1. Index your documents - Create embeddings for all .txt files in the data/ folder:

python index.py

To force a specific device (useful for benchmarking):

python index.py --device cpu   # Force CPU
python index.py --device mps   # Force Apple Silicon GPU
python index.py --device cuda  # Force NVIDIA GPU

2. Search - Run the interactive search interface:

python search.py

You can also specify a device for search:

python search.py --device cpu

Type your query and the system will return the most similar documents based on semantic meaning.

How It Works

• index.py reads all .txt files from data/, computes their embeddings using the all-MiniLM-L6-v2 model, and saves them to embeddings.json
• search.py loads the embeddings, encodes your query, and finds the most similar documents using cosine similarity
• Sample data files are provided in data/ for testing
• The all-MiniLM-L6-v2 model produces 384-dimensional embedding vectors, providing a good balance between quality and efficiency
• Both scripts automatically detect and use GPU acceleration (Apple Silicon MPS or NVIDIA CUDA) when available, falling back to CPU if needed

Adding Your Own Data

Simply add .txt files to the data/ folder and re-run python index.py to update the embeddings.