Build the search function

Everything is in place: a corpus, its embeddings, and a way to score any query against them. Now we package it into the one function that is the search engine — embed the query, score it against the corpus, and return the best few matches.

Add this to search.py (replace the loose query-scoring lines from the last chapter):

import numpy as np

def search(query, k=3):
  query_embedding = model.encode(query, normalize_embeddings=True)
  scores = corpus_embeddings @ query_embedding      # cosine similarity, all docs

  top_k = np.argsort(-scores)[:k]                   # indices of the k highest scores
  return [(corpus[i], float(scores[i])) for i in top_k]

The only new piece is np.argsort(-scores): argsort returns the indices that would sort the scores ascending, so sorting the negated scores puts the highest first. Slice off the first k and you have your top matches.

Try it

Add a few searches to the bottom of search.py and run it:

for query in [
  "I can't log into my account",
  "when will my package arrive",
  "how do I get my money back",
]:
  print(f"\nQuery: {query}")
  for doc, score in search(query):
      print(f"  {score:.3f}  {doc}")

python search.py

Query: I can't log into my account
0.531  How do I reset my password?
0.334  How can I contact customer support?
0.286  How do I cancel my subscription?

Query: when will my package arrive
0.418  How do I update my shipping address?
0.377  Do you ship internationally?
0.331  Where can I see my past orders?

Query: how do I get my money back
0.604  What is your refund and return policy?
0.310  How do I cancel my subscription?
0.287  What payment methods do you accept?

Read those top hits. Not one query reuses the vocabulary of the document it matched — "money back" finds refund policy, "package arrive" finds shipping address — and the right answer lands first every time. That's a semantic search engine, and you built it in about a dozen lines.

Try your own queries. Add documents to the corpus list — anything you like — and search will index them the moment you re-run. The last chapter makes it something you'd actually ship.