Suppose you have an image but no filename, or a question but none of the keywords used in a document. An ID lookup or keyword search may miss what you need. Similarity search instead looks for nearby representations in a vector space: encode the query as an embedding, compare it with stored embeddings, and return nearest neighbors. Euclidean distance, cosine similarity, and dot product are possible comparison measures. Reverse image search is one practical use. section 3.4.5
For document retrieval, the workflow is:
- Break documents into chunks, embed them, and load those embeddings into a vector database. Use the embedding model planned for inference when preparing document embeddings. section 4.4.2 section 6.1.7
- Embed the incoming query and search for nearby document embeddings. In retrieval-augmented generation (RAG), put the retrieved text into the LLM prompt as context. section 6.1.7
- Budget for retrieved context. Three 400-token chunks add 1,200 tokens to a request; a coding assistant can instead select up to five examples and discard those that exceed its context window. section 6.1.7 section 10.2.3
Two uses of the same search step
A Wikipedia RAG pipeline chunks pages, embeds them, and uploads them to a Pinecone index so a query can retrieve relevant passages. A coding assistant searches a Milvus collection for examples near the user’s request, then formats the results as few-shot prompt examples. The retrieved neighbors serve different purposes, but both begin with an embedded query. In the coding-assistant test, the resulting answer was slightly worse than the previous one: adding retrieved examples does not guarantee a better response. section 6.1.7 section 10.2.3
Listing 4.11 Example embeddings
import numpy as np
from sentence_transformers import SentenceTransformer
from datasets import load_dataset
model_ckpt = "sentence-transformers/all-MiniLM-L6-v2" #1
model = SentenceTransformer(model_ckpt)
embs_train = load_dataset("tweet_eval", "emoji", split="train[:1000]")
embs_test = load_dataset("tweet_eval", "emoji", split="test[:100]")
def embed_text(example): #2
embedding = model.encode(example["text"])
return {"embedding": np.array(embedding, dtype=np.float32)}
print(f"Train 1: {embs_train[0]}")
embs_train = embs_train.map(embed_text, batched=False)
embs_test = embs_test.map(embed_text, batched=False)
embs_train.add_faiss_index("embedding") #3
#
idx, knn = 1, 3 # Select the first query and 3 nearest neighbors #4
query = np.array(embs_test[idx]["embedding"], dtype=np.float32)
scores, samples = embs_train.get_nearest_examples("embedding", query, k=knn)
print(f"QUERY LABEL: {embs_test[idx]['label']}") #5
print(f"QUERY TEXT: {embs_test[idx]['text'][:200]} [...]\n")
print("=" * 50)
print("Retrieved Documents:")
for score, label, text in zip(scores, samples["label"], samples["text"]):
print("=" * 50)
print(f"TEXT:\n{text[:200]} [...]")
print(f"SCORE: {score:.2f}")
print(f"LABEL: {label}")
A small Python example embeds tweets with SentenceTransformer, indexes their embeddings with FAISS, and selects three nearest neighbors for a test query. section 4.4.2
Choose the search setup for the work around it, not just for vector storage. A vector database provides fast retrieval; a feature store can sit on top of one to help manage multiple data sources, embedding-model experiments, and frequent updates. Similarity also needs to suit the comparison goal: finding nearby vectors is not, by itself, a guarantee that the results will be useful. section 3.4.5 section 4.4.2
