Now that we have made some candidate semantic chunks, it might be useful to see how similar they are to one another. This will help us get a sense for what information they contain. We will proceed by embedding the semantic chunks, and then use UMAP to reduce the dimensionality of the resulting embeddings to 2D so that we can plot them.
UMAP stands for Uniform Manifold Approximation and Projection, and is a powerful, general dimensionality reduction technique that can capture non-linear relationships. A full explanation of how it works can be found here. The purpose of using it here is to capture something of the relationships that exist between the embedded chunks in 1536-D space in a 2-D plot
from umap import UMAPdimension_reducer = UMAP(
n_neighbors=5,
n_components=2,
min_dist=0.0,
metric="cosine",
random_state=0
)
reduced_embeddings = dimension_reducer.fit_transform(semantic_embeddings)
splits_df = pd.DataFrame(
{
"reduced_embeddings_x": reduced_embeddings[:, 0],
"reduced_embeddings_y": reduced_embeddings[:, 1],
"idx": np.arange(len(reduced_embeddings[:, 0])),
}
)
splits_df["chunk_end"] = np.cumsum([len(x) for x in semantic_text_groups])
ax = splits_df.plot.scatter(
x="reduced_embeddings_x",
y="reduced_embeddings_y",
c="idx",
cmap="viridis"
)
ax.plot(
reduced_embeddings[:, 0],
reduced_embeddings[:, 1],
"r-",
linewidth=0.5,
alpha=0.5,
)
UMAP is quite sensitive to the n_neighbors parameter. Generally the smaller the value of n_neighbors, the more the algorithm focuses on the use of local structure to learn how to project the data into lower dimensions. Setting this value too small can lead to projections that don’t do a great job of capturing the large scale structure of the data, and it should generally increase as the number of datapoints grows.
A projection of our data is shown below and its quite informative: Clearly we have three clusters of similar meaning, with the 1st and 3rd being more similar to each other than either is to the 2nd. The idx color bar in the plot above shows the chunk number, while the red line gives us an indication of the sequence of the chunks.
What about automatic clustering? This would be helpful if we wanted to group the chunks into larger segments or topics, which could serve as useful metadata to filter on in a RAG application with hybrid search, for example. We also might be able to group chunks that are far apart in the text (and therefore would not have been grouped by the standard semantic chunking in section 1) but have similar meanings.
There are many clustering approaches that could be used here. HDBSCAN is a possibility, and is the default method recommended by the BERTopic package. However, in this case hierarchical clustering seems more useful since it can give us a sense of the relative importance of whatever groups emerge. To run hierarchical clustering, we first use UMAP to reduce the dimensionality of the dataset to a smaller number of components. So long as UMAP is working well here, the exact number of components shouldn’t significantly affect the clusters that get generated. Then we use the hierarchy module from scipy to perform the clustering and plot the result using seaborn
from scipy.cluster import hierarchy
from scipy.spatial.distance import pdist
from umap import UMAP
import seaborn as sns# set up the UMAP
dimension_reducer_clustering = UMAP(
n_neighbors=umap_neighbors,
n_components=n_components_reduced,
min_dist=0.0,
metric="cosine",
random_state=0
)
reduced_embeddings_clustering = dimension_reducer_clustering.fit_transform(
semantic_group_embeddings
)
# create the hierarchy
row_linkage = hierarchy.linkage(
pdist(reduced_embeddings_clustering),
method="average",
optimal_ordering=True,
)
# plot the heatmap and dendogram
g = sns.clustermap(
pd.DataFrame(reduced_embeddings_clustering),
row_linkage=row_linkage,
row_cluster=True,
col_cluster=False,
annot=True,
linewidth=0.5,
annot_kws={"size": 8, "color": "white"},
cbar_pos=None,
dendrogram_ratio=0.5
)
g.ax_heatmap.set_yticklabels(
g.ax_heatmap.get_yticklabels(), rotation=0, size=8
)
The result is also quite informative. Here n_components_reduced was 4, so we reduced the dimensionality of the embeddings to 4D, therefore making a matrix with 4 features where each row represents one of the semantic chunks. Hierarchical clustering has identified the two major groups (i.e. trees and Namibia), two large subgroup within trees (i.e. medical uses vs. other) and an number of other groups that might be worth exploring.
Note that BERTopic uses a similar technique for topic visualization, which could be seen as an extension of what’s being presented here.
How is this useful in our exploration of semantic chunking? Depending on the results, we may choose to group some of the chunks together. This is again quite subjective and it might be important to try out a few different types of grouping. Let’s say we looked at the dendrogram and decided we wanted 8 distinct groups. We could then cut the hierarchy accordingly, return the cluster labels associated with each group and plot them.
cluster_labels = hierarchy.cut_tree(linkage, n_clusters=n_clusters).ravel()
dimension_reducer = UMAP(
n_neighbors=umap_neighbors,
n_components=2,
min_dist=0.0,
metric="cosine",
random_state=0
)
reduced_embeddings = dimension_reducer.fit_transform(semantic_embeddings)splits_df = pd.DataFrame(
{
"reduced_embeddings_x": reduced_embeddings[:, 0],
"reduced_embeddings_y": reduced_embeddings[:, 1],
"cluster_label": cluster_labels,
}
)
splits_df["chunk_end"] = np.cumsum(
[len(x) for x in semantic_text_groups]
).reshape(-1, 1)
ax = splits_df.plot.scatter(
x="reduced_embeddings_x",
y="reduced_embeddings_y",
c="cluster_label",
cmap="rainbow",
)
ax.plot(
reduced_embeddings[:, 0],
reduced_embeddings[:, 1],
"r-",
linewidth=0.5,
alpha=0.5,
)
The resulting plot is shown below. We have 8 clusters, and their distribution in the 2D space looks reasonable. This again demonstrates the importance of visualization: Depending on the text, application and stakeholders, the right number and distribution of groups will likely be different and the only way to check what the algorithm is doing is by plotting graphs like this.
Assume after a few iterations of the steps above, we’ve settled on semantic splits and clusters that we’re happy with. It then makes sense to ask what these clusters actually represent? Obviously we could read the text and find out, but for a large corpus this is impractical. Instead, let’s use an LLM to help. Specifically, we will feed the text associated with each cluster to GPT-4o-mini and ask it to generate a summary. This is a relatively simple task with LangChain, and the core aspects of the code are shown below
import langchain
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers.string import StrOutputParser
from langchain.callbacks import get_openai_callback
from dataclasses import dataclass@dataclass
class ChunkSummaryPrompt:
system_prompt: str = """
You are an expert at summarization and information extraction from text. You will be given a chunk of text from a document and your
task is to summarize what's happening in this chunk using fewer than 10 words.
Read through the entire chunk first and think carefully about the main points. Then produce your summary.
Chunk to summarize: {current_chunk}
"""
prompt: langchain.prompts.PromptTemplate = PromptTemplate(
input_variables=["current_chunk"],
template=system_prompt,
)
class ChunkSummarizer(object):
def __init__(self, llm):
self.prompt = ChunkSummaryPrompt()
self.llm = llm
self.chain = self._set_up_chain()
def _set_up_chain(self):
return self.prompt.prompt | self.llm | StrOutputParser()
def run_and_count_tokens(self, input_dict):
with get_openai_callback() as cb:
result = self.chain.invoke(input_dict)
return result, cb
llm_model = "gpt-4o-mini"
llm = ChatOpenAI(model=llm_model, temperature=0, api_key=api_key)
summarizer = ChunkSummarizer(llm)
Running this on our 8 clusters and plotting the result with datamapplot gives the following