Semantic search: the k nearest

Code in:JavaC#PythonJavaScript

In short: Semantic search compares the vector of the question with the vectors of the texts and returns the k nearest by cosine, even with no word in common.

Problem

One answer is not enough

The answer base of the support team now has six texts:

// the answer base of the support team
const DOCS = [
  { topic: "delivery", text: "delivery takes up to two days" },
  { topic: "returns", text: "return the item within 14 days and get the money" },
  { topic: "payment", text: "pay by card or by bank transfer" },
  { topic: "price", text: "delivery costs 5 euro to an address and 3 euro to an office" },
  { topic: "declined", text: "if the bank says no try another card or pay the courier" },
  { topic: "warranty", text: "the warranty is two years from the purchase" },
];

// the questions of the customers
const QUESTIONS = [
  "when will my parcel come",
  "i want to return my order",
  "how much is delivery to an office",
];

The bot from the lesson on vectors turns every text into a vector and returns only the one nearest to the question:

Output

when will my parcel come → delivery: delivery takes up to two days
i want to return my order → delivery: delivery takes up to two days
how much is delivery to an office → price: delivery costs 5 euro to an address and 3 euro to an office

The first and the third question are fine. The second customer wants to send the order back and gets the delivery time. The returns text is not far: it is second, 0.788 against 0.853. But the bot looks only at the first one.

The idea

The k nearest

The search should not bet everything on one text. Let it return the first k: sorted from the nearest, together with the similarity.

This is semantic search:

1. Every text in the base becomes a vector. The vectors are computed once and kept next to the texts. This is the index.
2. The question becomes a vector with the same space.
3. The cosine of the question to every entry in the index.
4. Sorting from the largest cosine and the first k.

These are the nearest neighbours from the lesson on k-NN. There the neighbours were past orders, here they are texts. The measure is the cosine, not the distance.

The k found texts are not the answer. They go to the language model, which reads them and answers from them. That is the next step.

See it

The index

The vectors come from the same space as in the lesson on vectors: the eight most frequent words in the shop text. The index is a list of entries: the text and its vector.

// the index: the vector of every text, computed once
function buildIndex(space) {
  return DOCS.map((doc) => ({ doc, vector: space.ofText(doc.text) }));
}

const space = new Space(TEXT, STOP, 8);
const index = buildIndex(space);
for (const e of index)
  console.log(e.doc.topic.padEnd(10) + `[${e.vector.join(", ")}]`);

Output

delivery  [2, 2, 2, 1, 2, 0, 2, 1]
returns   [0, 6, 2, 1, 4, 1, 3, 1]
payment   [0, 0, 0, 0, 0, 4, 0, 0]
price     [9, 1, 12, 2, 1, 0, 1, 0]
declined  [0, 0, 2, 2, 0, 5, 0, 1]
warranty  [0, 0, 0, 0, 0, 0, 0, 0]

"payment" has only the axis "card". "warranty" is all zeros: its words are not in the shop text, and the small model does not know what it is about. We come back to this in "When it hurts".

See it

The first k

The search takes the vector of the question, computes the cosine to every entry, sorts from the nearest and returns the first k:

// the first k texts nearest to the question vector
function search(index, v, k) {
  const hits = [];
  for (const e of index)
    hits.push({ doc: e.doc, score: cosine(v, e.vector) });
  hits.sort((a, b) => b.score - a.score);
  return hits.slice(0, k);
}

for (const q of QUESTIONS) {
  console.log(q);
  const v = space.ofText(q);
  for (const h of search(index, v, 3)) {
    const topic = h.doc.topic.padEnd(10);
    console.log(`  ${topic}${h.score.toFixed(3)}`);
  }
}

Output

when will my parcel come
  delivery  0.696
  price     0.536
  returns   0.347
i want to return my order
  delivery  0.853
  returns   0.788
  price     0.394
how much is delivery to an office
  price     0.763
  delivery  0.715
  returns   0.353

For the return the first text is still "delivery", but the second one is "returns". With k = 3 it is among the hits. Here is the same as a picture. The vectors have 8 dimensions, so the picture keeps only how far each text is from the question:

// the first k texts nearest to the question vector
function search(index, v, k) {
  const hits = [];
  for (const e of index)
    hits.push({ doc: e.doc, score: cosine(v, e.vector) });
  hits.sort((a, b) => b.score - a.score);
  return hits.slice(0, k);
}

const v = space.ofText("i want to return my order");
for (const h of search(index, v, 3)) {
  const topic = h.doc.topic.padEnd(10);
  console.log(`  ${topic}${h.score.toFixed(3)}`);
}

Step by step, in words

1. The question "i want to return my order" becomes a vector: [2, 2, 0, 0, 2, 0, 2, 0]. It is the point in the centre.
2. The cosine to every entry in the index. The bigger the cosine, the closer the point. A text with no axis in common is on the outer circle.
3. Sorted from the nearest: "delivery", "returns", "price"…
4. k = 1: only "delivery" (0.853). The customer wants to send the order back and gets the delivery time.
5. k = 2: "returns" (0.788) is in. The right text is among the hits.
6. k = 3: "price" (0.394) too. More texts: more chance the right one is in, and more noise.
7. The k = 3 nearest: "delivery", "returns", "price". The right text, "returns", is second, so the search returns k texts, not one.

Try it

Predict the output

The customer writes "can i change the colour". None of the words is in the shop text. What does the program print?

const v = space.ofText("can i change the colour");
for (const h of search(index, v, 3))
  console.log(`${h.doc.topic} ${h.score.toFixed(3)}`);

Answer

delivery 0.000
returns 0.000
payment 0.000

price 0.000
declined 0.000
warranty 0.000

warranty 0.000

Check

Try it

Fill in

The search has to sort from the largest similarity to the smallest. What goes in the place of ___?

// the first k texts nearest to the question vector
function search(index, v, k) {
  const hits = [];
  for (const e of index)
    hits.push({ doc: e.doc, score: cosine(v, e.vector) });
  hits.sort((a, b) => );
  return hits.slice(0, k);
}

Check

Try it

Find the bug

A colleague writes the search, and for "i want to return my order" gets "payment", "declined" and "warranty": all texts with similarity 0.000. On which line is the bug?

Pick the line with the bug.

1// a colleague returns the first k2function search(index, v, k) {3 const hits = [];4 for (const e of index)5 hits.push({ doc: e.doc, score: cosine(v, e.vector) });6 hits.sort((a, b) => a.score - b.score);7 return hits.slice(0, k);8}

Check

The idea

How big should k be

With k = 1 the bot misses the return. With k = 3 the right text is in, but others come too: for "when will my parcel come" there are also "price" (0.536) and "returns" (0.347), which do not answer the question.

• A small k misses the right text when it is not first.
• A big k brings noise. Every text goes into the prompt: these are tokens that cost money and room in the context window. And the model can hold on to the wrong text.
• A threshold. A text below a set similarity is not returned at all. So a question like "can i change the colour" does not get three random texts, but "not sure".

Real systems usually take between 3 and 10 texts, but that is only an example. How many exactly is measured with a test set of questions for which you know the right text. That is evaluating the search.

The idea

When it hurts

• Exact things. "where is order 10234": the number is an unknown word, and its vector is zeros. Vectors catch meaning, not exact strings. An order number, a product code or an email is looked up by equality in the database. For names and rare words a search by words helps, and real systems often mix the two. Search engines over words are in the lesson on searching text.
• Words the model does not know. "warranty" is 0.000 to every question and will never be found. A real model knows many more words, but the same happens with the company's own jargon and with new product names.
• Every question goes through the whole index. For six texts that is six cosines. For a million texts with vectors of 1024 numbers each it is about a billion multiplications per question (an example number). That is why vector databases search approximately: they check a small part of the entries and almost always find the same k. They are in the lesson on vector databases. Running and maintaining such a database is a job for ZeroToOps.
• The index gets old. Change a text, and its vector is computed again. Change the model for vectors, and the whole index is computed again: vectors from two models are not compared with each other.
• A long document, one vector. The vector of a whole page is a mix of all its topics. How a document is cut into pieces is in the next lesson.

More precisely, if you are curious

The real tools. The vectors come from a model for vectors: it is called through an API like the other models, or run locally, most often from Python. They are kept in a vector database (for example pgvector in PostgreSQL), which has an index for approximate search. In Python the computing with vectors is usually done with NumPy. ML practice is mostly in Python, but the search is called from any language: the query to the database is the same. Here everything is by hand, so every number is visible.

Sorting everything or only the first k. Sorting n entries is O(n log n). If k is small, a heap of k elements is enough: every entry is compared with the weakest one in it, and that is O(n log k) steps. The heap is in the lesson on the priority queue.

Equal similarities. The sort in all four languages is stable: equal values stay in the order they have in the base. That is why "predict" prints the first three texts. Do not rely on this for the meaning: equal similarities mean that the model cannot tell the texts apart.

The distance in the picture. The point of every text is √(2 − 2·cos) away from the question. That is the distance between the tips of the two vectors if both have length 1. Sorting by it gives the same order as sorting by the cosine.

Exercise

At most k, but only the close ones

Solve it in your editor.

Write a function that returns at most k texts, but only those with similarity at least 0.5. If none is left, it returns "not sure, passing it to a colleague". Try it with k = 3 for "when will my parcel come", "i want to return my order" and "can i change the colour".

Output

when will my parcel come: delivery 0.696, price 0.536
i want to return my order: delivery 0.853, returns 0.788
can i change the colour: not sure, passing it to a colleague

The first question keeps "delivery" and "price": "returns" is below the threshold. The second keeps "delivery" and "returns", and the third keeps nothing and goes to a person.

More tasks

• A heap instead of a sort. Make the search with a heap of k elements. The output has to be the same.
• The gap between the first and the second. Print the difference in similarity between the first and the second text. A small gap means the search is not sure.

Test

Check yourself

1. Why does the search return k texts and not only the nearest one?

Because the search is faster that wayBecause the language model only takes three textsThe nearest one is not always the right one, and with k texts the right one is more often among them

2. What does the index keep?

The vector of every text in the base, computed onceThe ready answers of the modelThe questions of the customers

3. The question has not a single word the model knows. What does the search return with k = 3?

An errorThree texts with similarity 0.000Nothing

4. The customer writes "where is order 10234". Why does semantic search not help here?

The number is an exact identifier: the vectors do not know it, and the order is looked up directly in the database by numberBecause the base has no text about deliveryBecause the question is too short

Check

Start the lesson over

Related lessons

• Builds on: What machine learning is · Text in pieces: tokens · Predicting the next token · Embeddings: meaning as vectors
• Continues in: Chunking documents · RAG: answers from the documents · Evaluating RAG: does it find it, does it stand on it · The support assistant
• Tasks on this topic: Semantic search and RAG