Opening the Black Box: Explaining Text Embeddings for Semantic Search
You might have heard about techniques such as ”semantic search”, ”vector search”, or ”embedding search” becoming more popular in recent years. This “new” way of searching allows a user to find results that do not exactly match the words the user typed in the searchbar, but also get results that are similar in meaning. For example, when searching for “front row seat” you will be interested in “first tier tickets”, and not in an article named “Rowing: differences between sitting in the front and back seat”. You can find an introduction to semantic search and its comparison to regular keyword search in this blog:
The reason these techniques have become popular in recent years is due to a significant improvement in the understanding of linguistic-context by language models. Researchers and AI-companies are doing their best to train the best embedding models (or Large Language Models) for different tasks. Better LLM’s also lead to better embedding models, because the first layer of an LLM does precisely that: contextualize the input text.
Discussing all use cases for text embeddings or diving deep into how LLM’s work are topics on their own, instead I will zoom in on embedding search, what it is, and how to explain its search results. I will also provide code which makes any embedding model a “glass box”: using the inner workings of the model to explain itself.
Why explain embeddings?
Explaining embeddings is not just a technical curiosity; it is a necessity. Users of AI-search engines want intuitive results, something you only get by explaining the underlying embedding. Moreover, with the EU AI Act emphasizing accountability and compliance, it might also be required that the output of a search enginge can be explained in a clear way, see article 14 of the AI Act.
So wether you wish to improve the user experience of your search engine, or you don’t want to be an AI-cowboy running from the law, you need an explainable search system.
Text embeddings
I’ve been talking about them for a while: text embeddings. But what are they actually? To answer this, I’ll first explain how a model reads input-text, before explaining what the model does to arrive at the meaning of the text and the text embedding.
Embedding Models and their big cousin Large Language Models are neural networks. A neural network cannot read text, so instead it has to translate text into numbers. These numbers are called tokens and the part of the model that does this translation is called the tokenizer.
After the tokenizer step, all further steps opperate on tokens. Not every word has its own token, some words are split into mulitple tokens and some words bring “special tokens” with them. Here are some examples of words and their tokens:
The model has a “meaning” stored for each token in the form of a vector: a collection of numbers. By placing all these vectors in a matrix, we get a large amount of numbers to represent the meaning of a text, which is good for a model… but not readable by a human. And we are not even done: the model still needs to combine these vectors to calculate the “meaning” of the text.
To do this, the model does two steps: it has learned how tokens are correlated to eachother in its attention-layer and how to combine the meanings of the tokens and their correlation in another layer of the neural network. The combination of the attention- and neural-network layer is called the transformer architecture, which is the backbone of all LLMs.
After applying these two steps to the vectors of each token, we once again have a vector for each token: the token embeddings. Each token embedding encapsulates the meaning of that token in the complete text. To get to the meaning of the whole text we do a last simple step: the model averages the vectors in the dimension of the tokens, which results in a single vector representing the meaning of the whole text: the text embedding. In embedding search, the user-query and stored text are both embedded and the similarity of the two text embeddings is calculated.
This is most likely too abstract for the users of your search engine to understand. But you still want to give them an intuitive explanation for the search results of your search engine. I will not showcase how to implement embedding search for a search engine. If you want to implement it yourself, I would suggest you look into LangChain’s vector stores, or ElasticSearch KNN search. Instead, I will provide you with code that returns an intuitive explanation: how much does each word correspond to the search query of a user?
In essence the code creates the embeddings for both a query and a search result and calculate the similariy, except for the search result you calculate the embedding of each word instead of for the complete text. The end result is how similar the meaning of each word in the search result is to the query. By using the inner workings of the AI model to explain itself, we get human-readable explanations of what words in the text caused the model to consider two texts to be similar in meaning.
The code
To start off, we need a model that can create embeddings. To illustrate the example I have chosen for the “small” model e5-base-v2 from intfloat, for which you can find the details about in their HuggingFace repository.
To find the latest models that are good at creating embeddings, you can look at the MTEB HuggingFace leaderboard, which shows how good the embeddings of many models are for different use cases of embeddings. To find the best models for semantic search: look at the “Retrieval” column.
import torch
import numpy as np
from transformers import pipeline, AutoModel, AutoTokenizer
from sentence_transformers.util import cos_sim
# example embedding model
model_name = "intfloat/e5-base-v2"
# load the model and tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)Depending on which model you use, the method for getting a sentence embedding from it is different. I have isolated the code for getting an embedding from the example model, you should change this when using a different model.
def embed(
input: str,
model: AutoModel,
tokenizer: AutoTokenizer,
context_length: int = 512
) -> torch.Tensor:
"""Calculate sentence embedding from embedding model.
See https://huggingface.co/intfloat/e5-base-v2#usage for the usage of the example model."""
encoding = tokenizer(input, return_tensors='pt', max_length=context_length, padding=True, truncation=True)
model_output = model(**encoding)
last_hidden = model_output.last_hidden_state.masked_fill(~encoding.attention_mask[..., None].bool(), 0.0)
return (last_hidden.sum(dim=1) / encoding.attention_mask.sum(dim=1)[..., None])Next we define two more functions: one that uses the pipeline class from transformers to get the token embeddings from the model, the other averages token embeddings per word to get “word embeddings” and then calculates the cosine similarity with the embedding of the original query.
def get_token_embeddings(
search_result: str,
model: AutoModel,
tokenizer: AutoTokenizer,
) -> torch.Tensor:
pipe = pipeline("feature-extraction", model=model, tokenizer=tokenizer)
token_embeddings = pipe(search_result)
return torch.Tensor(token_embeddings).squeeze()
def explain_search_result(
query: str,
search_result: str,
model: AutoModel,
tokenizer: AutoTokenizer,
context_length: int = 512,
) -> tuple[list[str], list[float]]:
"""Explains a search result by calculating the similarity of each word with the query.
Args:
query: The search term.
search_result: The search result to explain.
model: The embedding model.
tokenizer: The tokenizer of the embedding model.
context_length: Maximum context length of the model. Defaults to 512.
Returns:
tuple[list[str], list[float]]: A tuple of the list of all words and their cosine similarity with the query.
"""
# get token embeddings
result_encoded = tokenizer(search_result, return_tensors='pt', max_length=context_length, truncation=True)
query_embedding = embed(input=query, model=model, tokenizer=tokenizer, context_length=context_length)
query_embedding = query_embedding.detach().numpy()[0]
token_embeddings = get_token_embeddings(search_result, model, tokenizer)
# remove special tokens from token embeddings
word_ids = result_encoded.word_ids()
none_word_positions = [index for index, word in enumerate(word_ids) if word == None]
for index in sorted(none_word_positions, reverse=True):
del word_ids[index]
token_embeddings = torch.cat([token_embeddings[:index,], token_embeddings[index+1:]])
# calculate average word embedding from token embeddings
# https://stackoverflow.com/questions/56154604/groupby-aggregate-mean-in-pytorch
word_ids_tensor = torch.Tensor(word_ids).int()
num_words = len(torch.unique(word_ids_tensor))
M = torch.zeros(num_words, len(token_embeddings))
M[word_ids_tensor, torch.arange(len(token_embeddings))] = 1
M = torch.nn.functional.normalize(M, p=1, dim=1)
word_embeddings = torch.mm(M, token_embeddings)
# calculate similarity per word and the input query
word_query_similarities = cos_sim(query_embedding, word_embeddings)[0]
# get character spans from original text
char_spans = [result_encoded.word_to_chars(word_id) for word_id in set(word_ids)]
return ([search_result[char_spans[index][0]:char_spans[index][1]]
for index in range(len(char_spans))
],
[float(similarity) for similarity in word_query_similarities])Example usage
Here is an example on how to use the functions defined above. This query and search results will also be the examples of the visualization.
# the example embedding model expects the input query to start with "query:"
# and the search results to start with "passage:"
query = "query: what is text embedding"
good_result = """passage:
In natural language processing (NLP), a word embedding is a representation of
a word. The embedding is used in text analysis. Typically, the representation
is a real-valued vector that encodes the meaning of the word in such a way
that the words that are closer in the vector space are expected to be similar
in meaning. Word embeddings can be obtained using language modeling and feature
learning techniques, where words or phrases from the vocabulary are mapped to
vectors of real numbers."""
bad_result = """passage:
An embedded system is a computer system—a combination of a computer
processor, computer memory, and input/output peripheral
devices—that has a dedicated function within a larger mechanical
or electronic system. It is embedded as part of a complete device
often including electrical or electronic hardware and mechanical
parts. Because an embedded system typically controls physical
operations of the machine that it is embedded within, it often has
real-time computing constraints. Embedded systems control many
devices in common use. In 2009, it was estimated that ninety-eight
percent of all microprocessors manufactured were used in embedded
systems."""
# explain both results
good_result_explained = explain_search_result(query, good_result, model, tokenizer)
bad_result_explained = explain_search_result(query, bad_result, model, tokenizer)
# print both results and average similarity
print(good_result_explained)
print(sum(good_result_explained[1]) / len(good_result_explained[1]))
print(bad_result_explained)
print(sum(bad_result_explained[1]) / len(bad_result_explained[1]))Visualization
The whole idea of explaining semantic search results is to be able to show users why this result was returned, even though some of the words do not match exactly.
Below you will find code that writes the query and search result to an image, where less relevant words are light grey and most relevant words are dark grey or black.
from PIL import Image, ImageDraw, ImageFont
def visualize_explained_search_result(
search_query: str,
words: list[str],
scores: list[float],
relative_score: bool = True,
min_opacity: float = 0.25,
img_width: int = 400,
font_size: int = 20,
font_name: str = "arial.ttf",
line_margin: int = 4,
) -> Image:
"""Create an image for an explained search result. The image shows the query and search result
text. The colour of the search result is darker for words that are more similar to the
input query.
"""
img = Image.new("RGB", (img_width, len(words) * font_size), (255, 255, 255))
font = ImageFont.truetype(font_name, font_size)
draw = ImageDraw.Draw(img)
draw.text((0, 0), search_query, fill='black', font=font)
width = 0
start_height = 2 * font_size
height = start_height
min_similarity = 0
max_similarity = 1
if relative_score:
min_similarity = min(scores)
max_similarity = max(scores)
for word, similarity in zip(words, scores):
text = " " + word
length = font.getlength(text)
if width + length + line_margin > img_width:
width = 0
height += font_size + 4
sim = max(min_opacity, (similarity - min_similarity) / (max_similarity - min_similarity))
grey_color = round(255 * (1-sim))
draw.text((width, height), text, font=font, fill=(grey_color, grey_color, grey_color))
width += length
return img.crop((0, 0, img_width, height+start_height))
visualize_explained_search_result(query, good_result_explained[0], good_result_explained[1], False)
visualize_explained_search_result(query, bad_result_explained[0], bad_result_explained[1], False)
Conclusion
In conclusion, it is possible to explain search results from embedding search. By using the inner working of the embedding model, you are certain that your explanations align with the actual working of your model. The explanation show how much impact each word has on the semantic similarity of two texts. Explaining search results is important for transparancy to users and for compliancy with the EU AI Act.
