How semantic search works, with a Python example

Semantic search finds results by what your query means instead of the exact words in it. It turns your query and every document into vectors with an embedding model, then ranks the documents by how closely their vectors point in the same direction as the query's.

That's how a search for "fantasy movies" can return Lord of the Rings and Harry Potter, even though neither title has the word "fantasy" in it. Below I go through keyword search, vectors, embedding models and cosine similarity, then the full pipeline, and then a working movie search in Python. The video above covers the theory if you'd rather watch.

Keyword search vs semantic search

Keyword search matches words. Does this title or document contain the words in my query? If it does, it's a result. Search engines build an index from every word to the documents that contain it, so the lookup is fast and you can always see why something matched.

That works until the words don't line up. Search a movie database for "fantasy movies" with keywords and you'll get films with "fantasy" in the title or description. Lord of the Rings and Harry Potter don't say "fantasy" anywhere in their names, so you can't tell from the titles alone.

Semantic search matches the meaning behind the words. It knows "fantasy" has to do with magic and quests, and it knows Lord of the Rings is about a quest with wizards in it, so the two come out close together.

Keyword search still wins when someone types an exact product code or a person's name, which is why a lot of real search systems run both and merge the results (usually called hybrid search).

Words as vectors

Picture every word as a point on a map. "King" and "queen" mean similar things, so they sit close together. "Solar system" is a different concept entirely, so it sits far from both, near words like "Mars" and "Earth". Semantic search builds that map, and once you have it, "fantasy movies" lands near Lord of the Rings and Harry Potter.

Under the hood, those points are vectors: quantities with a magnitude and a direction, like the arrows from maths or physics class. Six feet tall is only a magnitude. A vector also points somewhere, and for search the direction is what matters. "King" and "queen" point almost the same way while "solar system" points somewhere else, so to check whether "fantasy" and "Lord of the Rings" are related, you check whether their vectors point the same way.

Real models use hundreds or a thousand numbers per vector, not two. The model I use below produces 1,024. You can't draw that, but the maths works the same as with arrows on paper.

Embedding models turn text into vectors

"Fantasy" is an English word, not a vector, so something has to convert it. That's the job of an embedding model: a machine learning model whose only purpose is to turn text or images into vectors.

The one I use most is thenlper/gte-large on Hugging Face. Hugging Face is a site full of AI models and datasets, and if you're doing anything with AI or machine learning, you should be on it.

You give the model text and it gives you back vectors, which are called embeddings. When someone says "embeddings", they just mean vectors that came out of an embedding model.

Cosine similarity ranks the results

Knowing two things are related isn't enough. If I search for "fantasy movies", should Lord of the Rings come first or Harry Potter? Obviously Lord of the Rings, because it's the better movie, but the search engine needs a number.

Because embeddings are numbers, you can run formulas on them, and the usual one is cosine similarity. It measures the angle between two vectors and ignores how long they are. Two vectors pointing the same way score 1, and the further apart they point, the lower the score. In theory it goes down to -1. With text embeddings the scores you see in practice almost all fall between 0 and 1, and higher means more similar.

So you score the query against every document and put the highest score at the top of the page. Euclidean distance (the straight-line distance between the two points) and the dot product also work, but cosine similarity is what I reach for in most cases.

The semantic search pipeline, step by step

Searching your own data by meaning takes two stages: you prepare the data once, then you answer queries as they come in.

Preparing the data:

1. Load your dataset, whether that's documents or rows from a database.
2. Convert it to text. Embedding models take text, so you can't pass in a database table as it is.
3. Clean it up and split it into chunks that fit the model's token limit.
4. Run every chunk through the embedding model.
5. Store each chunk's embedding next to the item it came from, so you never have to compute it again.

Answering a query:

1. Embed the user's query with the same model you used for the dataset. Vectors from two different models can't be compared, even when they're the same length.
2. Compute the cosine similarity between the query's embedding and every stored embedding.
3. Sort by that score and show the top 10 or 20 results.

Chunk your text to fit the token limit

Every embedding model has a maximum input length, measured in tokens. Some take 1,024 tokens and some take 512 or fewer. gte-large takes 512. Put in more than that and you get bad embeddings, or the model cuts the text off without telling you (that's what the sentence-transformers library does). Check the limit in the model's documentation and chunk your text to fit.

Store your embeddings

Generating embeddings is slow. I once generated about 50,000 of them for a project and it took around four hours. You never want to do that every time your program starts, so save them.

Embeddings are vectors, and there are databases built to store and search vectors: Pinecone, Chroma, pgvector (a Postgres extension) and Weaviate. For a small project, a file on disk is fine, and that's what the code below does.

Build a semantic movie search in Python

When I made the video I said I'd implement this in code in a later one. Here it is in writing instead: a movie search over the plots of 34,886 films, using the pipeline above. The same code is in my semantic-search repo on GitHub, with setup steps for uv.

The data is the Wikipedia Movie Plots dataset on Kaggle. Download it and put wiki_movie_plots_deduped.csv in a data folder.

Install sentence-transformers, which pulls in PyTorch and downloads models from Hugging Face for you:

pip install sentence-transformers

If you're on Linux without an NVIDIA GPU, install the CPU build of PyTorch first with pip install torch --index-url https://download.pytorch.org/whl/cpu. The default one bundles a few gigabytes of CUDA libraries you won't use.

The project looks like this:

semantic-search/
├── main.py
├── data/
│   └── wiki_movie_plots_deduped.csv
└── cache/
    └── embeddings.pkl    (created on the first run)

Here's main.py:

import csv
import os
import pickle
import re
from dataclasses import dataclass

import torch
from sentence_transformers import SentenceTransformer

MODEL_NAME = "thenlper/gte-large"
# MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
DATASET_PATH = "data/wiki_movie_plots_deduped.csv"
CACHE_PATH = "cache/embeddings.pkl"


@dataclass
class Movie:
    release_year: int
    title: str
    origin: str
    director: str
    cast: str
    genre: str
    plot: str


def load_model() -> SentenceTransformer:
    # Half precision is fast on a GPU but very slow on a CPU
    if torch.cuda.is_available():
        return SentenceTransformer(MODEL_NAME, device="cuda", model_kwargs={"dtype": torch.float16})
    return SentenceTransformer(MODEL_NAME, device="cpu", model_kwargs={"dtype": torch.float32})


def load_dataset(file_path: str = DATASET_PATH) -> list[Movie]:
    movies = []
    with open(file_path, encoding="utf-8") as csvfile:
        for row in csv.DictReader(csvfile):
            movies.append(Movie(
                release_year=int(row["Release Year"]),
                title=row["Title"],
                origin=row["Origin/Ethnicity"],
                director=row["Director"],
                cast=row["Cast"],
                genre=row["Genre"],
                plot=row["Plot"],
            ))
    return movies


def divide_into_chunks(text: str, chunk_size: int = 3, overlap: int = 1) -> list[str]:
    sentences = [s.strip() for s in re.split(r"(?<=[.!?])\s+", text) if s.strip()]
    chunks = []
    step = chunk_size - overlap

    for i in range(0, len(sentences), step):
        chunks.append(" ".join(sentences[i : i + chunk_size]))
        if i + chunk_size >= len(sentences):
            break

    return chunks


def convert_to_text(movie: Movie) -> list[str]:
    # Every chunk repeats the movie's details, so each one makes sense on its own
    chunks = divide_into_chunks(movie.plot) or [""]
    return [
        f"Title: {movie.title}\n"
        f"Year: {movie.release_year}\n"
        f"Director: {movie.director}\n"
        f"Cast: {movie.cast}\n"
        f"Genre: {movie.genre}\n"
        f"Origin: {movie.origin}\n\n"
        f"Plot ({i + 1}):\n"
        f"{chunk}"
        for i, chunk in enumerate(chunks)
    ]


def load_or_create_embeddings(movies: list[Movie], model: SentenceTransformer) -> dict:
    if os.path.exists(CACHE_PATH):
        print("Loading embeddings from cache...")
        with open(CACHE_PATH, "rb") as f:
            return pickle.load(f)

    texts = []
    metadata = []
    for movie in movies:
        for i, text in enumerate(convert_to_text(movie)):
            texts.append(text)
            metadata.append({
                "title": movie.title,
                "year": movie.release_year,
                "genre": movie.genre,
                "director": movie.director,
                "chunk_index": i,
            })

    print(f"Embedding {len(texts)} chunks from {len(movies)} movies...")
    embeddings = model.encode(texts, batch_size=64, show_progress_bar=True, convert_to_numpy=True)

    data = {"embeddings": embeddings, "texts": texts, "metadata": metadata}
    os.makedirs(os.path.dirname(CACHE_PATH), exist_ok=True)
    with open(CACHE_PATH, "wb") as f:
        pickle.dump(data, f)

    return data


def search(query: str, model: SentenceTransformer, embedding_data: dict, top_k: int = 200) -> list[dict]:
    # The query goes through the same model as the dataset
    query_embedding = model.encode(query, convert_to_tensor=True)
    corpus_embeddings = torch.tensor(embedding_data["embeddings"], device=query_embedding.device)

    scores = torch.cosine_similarity(query_embedding, corpus_embeddings, dim=1)
    top_results = torch.topk(scores, k=min(top_k, len(scores)))

    results = []
    for score, idx in zip(top_results.values, top_results.indices):
        idx = idx.item()
        results.append({
            "score": score.item(),
            "text": embedding_data["texts"][idx],
            "metadata": embedding_data["metadata"][idx],
        })
    return results


def display_search_results(results: list[dict], limit: int = 20):
    # One movie can match with several chunks, so show each movie once
    seen = set()
    for result in results:
        meta = result["metadata"]
        if (meta["title"], meta["year"]) in seen:
            continue
        seen.add((meta["title"], meta["year"]))

        plot = result["text"].split(":\n", maxsplit=1)[-1].replace("\n", " ")
        if len(plot) > 200:
            plot = plot[:200] + "..."

        print(f"{len(seen)}) {meta['title']} ({result['score']:.4f}) [{meta['genre']}, {meta['year']}, {meta['director']}]")
        print(plot)
        print()

        if len(seen) >= limit:
            break


def main():
    movies = load_dataset()
    print(f"Loaded {len(movies)} movies.")

    model = load_model()
    embedding_data = load_or_create_embeddings(movies, model)

    while query := input("Search (blank line to quit): ").strip():
        display_search_results(search(query, model, embedding_data))


if __name__ == "__main__":
    main()

load_model() loads gte-large on the GPU if you have one. On a CPU it asks for full 32-bit precision, because the current version of sentence-transformers loads this model in half precision, which runs many times slower on a CPU.

divide_into_chunks() splits a plot into sentences and groups them three at a time, with one sentence of overlap so the context carries from one chunk into the next. convert_to_text() puts the movie's title, year, director, cast and genre on top of every chunk. So a chunk from the middle of a plot still knows which movie it belongs to.

load_or_create_embeddings() is the preparation stage. If cache/embeddings.pkl exists, it loads that. If not, it embeds every chunk in batches of 64 and pickles the vectors to disk along with each chunk's movie details.

search() is the query stage. It embeds the query with the same model and scores it against every chunk with torch.cosine_similarity. Then torch.topk picks the 200 best, which is faster than sorting the whole list. A long plot gets many chunks, so one movie can take several of those 200 spots. display_search_results() skips repeats and prints up to 20 different movies.

Run it

python main.py

The first run downloads gte-large (about 670 MB), then embeds every chunk and saves the result to cache/embeddings.pkl. After that, it loads the file and goes straight to the search prompt.

That first run is slow without a GPU. The whole dataset comes to 328,648 chunks, and on my laptop's CPU gte-large embedded fewer than two chunks a second, which works out to more than two days. On a machine with an NVIDIA GPU the code uses the GPU on its own, and Google Colab has free GPU runtimes if you don't have one. On a CPU, switch to the commented-out all-MiniLM-L6-v2 line. It's a much smaller model (384 numbers per vector and a 256-token limit) that ran about 20 times faster in my test, at some cost to the results. I tested this code on 350 of the movies (4,868 chunks), and with gte-large that took 46 minutes on the same CPU.

Delete cache/embeddings.pkl whenever you change the model or the data, or you'll be searching stale vectors.

Here's the top of what I got for the query from the video (I've cut the plot line under each result):

Search (blank line to quit): fantasy movies
1) The Lord of the Rings: The Two Towers (0.8742) [adventure, fantasy, 2002, Peter Jackson]
2) The Chronicles of Narnia: Prince Caspian (0.8737) [fantasy, 2008, Andrew Adamson]
3) The Lord of the Rings: The Return of the King (0.8718) [adventure, fantasy, 2003, Peter Jackson]
4) The Lord of the Rings: The Fellowship of the Ring (0.8717) [fantasy, 2001, Peter Jackson]
5) The Sorcerer's Apprentice (0.8677) [adventure, comedy, fantasy, 2010, Jon Turteltaub]
6) The Lord of the Rings (0.8671) [fantasy, animated, 1978, Ralph Bakshi]

Lord of the Rings on top, as it should be. That query only printed 15 movies, because the Lord of the Rings and Harry Potter films have long plots and filled most of the 200 best chunks between them.

The genre field helps with that query, since it literally says "fantasy". Queries that don't name a genre show the difference better. "A wizard school" put Harry Potter and the Sorcerer's Stone first. "A man stranded alone on another planet" put The Martian first, and Cast Away and Interstellar both made the top four.

Further down that list the results drift into movies that have little to do with the query. Semantic search always returns its closest matches, even when nothing is close, and a 350-movie slice doesn't have many films about being stranded. The score tells you how close each match is, so in a real app you'd cut results off below a threshold.

Why this is worth learning

This is how Google gives you decent results when your words don't match the page, and it's how RAG apps find the right documents before an LLM writes its answer. Under all of it are vectors and the angles between them.

I only learned this properly a month or two before I made the video, and I still find it fascinating. What looks like a search engine reading your mind is a model turning text into arrows and a formula checking which ones point the same way. Technology is beautiful once you understand how it works.