mirror of
https://github.com/hwchase17/langchain
synced 2024-11-11 19:11:02 +00:00
858ce264ef
```SemanticChunker``` currently provide three methods to split the texts semantically:
- percentile
- standard_deviation
- interquartile
I propose new method ```gradient```. In this method, the gradient of distance is used to split chunks along with the percentile method (technically) . This method is useful when chunks are highly correlated with each other or specific to a domain e.g. legal or medical. The idea is to apply anomaly detection on gradient array so that the distribution become wider and easy to identify boundaries in highly semantic data.
I have tested this merge on a set of 10 domain specific documents (mostly legal).
Details :
- **Issue:** Improvement
- **Dependencies:** NA
- **Twitter handle:** [x.com/prajapat_ravi](https://x.com/prajapat_ravi)
@hwchase17
---------
Co-authored-by: Raviraj Prajapat <raviraj.prajapat@sirionlabs.com>
Co-authored-by: isaac hershenson <ihershenson@hmc.edu>
283 lines
10 KiB
Python
283 lines
10 KiB
Python
"""Experimental **text splitter** based on semantic similarity."""
|
|
import copy
|
|
import re
|
|
from typing import Any, Dict, Iterable, List, Literal, Optional, Sequence, Tuple, cast
|
|
|
|
import numpy as np
|
|
from langchain_community.utils.math import (
|
|
cosine_similarity,
|
|
)
|
|
from langchain_core.documents import BaseDocumentTransformer, Document
|
|
from langchain_core.embeddings import Embeddings
|
|
|
|
|
|
def combine_sentences(sentences: List[dict], buffer_size: int = 1) -> List[dict]:
|
|
"""Combine sentences based on buffer size.
|
|
|
|
Args:
|
|
sentences: List of sentences to combine.
|
|
buffer_size: Number of sentences to combine. Defaults to 1.
|
|
|
|
Returns:
|
|
List of sentences with combined sentences.
|
|
"""
|
|
|
|
# Go through each sentence dict
|
|
for i in range(len(sentences)):
|
|
# Create a string that will hold the sentences which are joined
|
|
combined_sentence = ""
|
|
|
|
# Add sentences before the current one, based on the buffer size.
|
|
for j in range(i - buffer_size, i):
|
|
# Check if the index j is not negative
|
|
# (to avoid index out of range like on the first one)
|
|
if j >= 0:
|
|
# Add the sentence at index j to the combined_sentence string
|
|
combined_sentence += sentences[j]["sentence"] + " "
|
|
|
|
# Add the current sentence
|
|
combined_sentence += sentences[i]["sentence"]
|
|
|
|
# Add sentences after the current one, based on the buffer size
|
|
for j in range(i + 1, i + 1 + buffer_size):
|
|
# Check if the index j is within the range of the sentences list
|
|
if j < len(sentences):
|
|
# Add the sentence at index j to the combined_sentence string
|
|
combined_sentence += " " + sentences[j]["sentence"]
|
|
|
|
# Then add the whole thing to your dict
|
|
# Store the combined sentence in the current sentence dict
|
|
sentences[i]["combined_sentence"] = combined_sentence
|
|
|
|
return sentences
|
|
|
|
|
|
def calculate_cosine_distances(sentences: List[dict]) -> Tuple[List[float], List[dict]]:
|
|
"""Calculate cosine distances between sentences.
|
|
|
|
Args:
|
|
sentences: List of sentences to calculate distances for.
|
|
|
|
Returns:
|
|
Tuple of distances and sentences.
|
|
"""
|
|
distances = []
|
|
for i in range(len(sentences) - 1):
|
|
embedding_current = sentences[i]["combined_sentence_embedding"]
|
|
embedding_next = sentences[i + 1]["combined_sentence_embedding"]
|
|
|
|
# Calculate cosine similarity
|
|
similarity = cosine_similarity([embedding_current], [embedding_next])[0][0]
|
|
|
|
# Convert to cosine distance
|
|
distance = 1 - similarity
|
|
|
|
# Append cosine distance to the list
|
|
distances.append(distance)
|
|
|
|
# Store distance in the dictionary
|
|
sentences[i]["distance_to_next"] = distance
|
|
|
|
# Optionally handle the last sentence
|
|
# sentences[-1]['distance_to_next'] = None # or a default value
|
|
|
|
return distances, sentences
|
|
|
|
|
|
BreakpointThresholdType = Literal[
|
|
"percentile", "standard_deviation", "interquartile", "gradient"
|
|
]
|
|
BREAKPOINT_DEFAULTS: Dict[BreakpointThresholdType, float] = {
|
|
"percentile": 95,
|
|
"standard_deviation": 3,
|
|
"interquartile": 1.5,
|
|
"gradient": 95,
|
|
}
|
|
|
|
|
|
class SemanticChunker(BaseDocumentTransformer):
|
|
"""Split the text based on semantic similarity.
|
|
|
|
Taken from Greg Kamradt's wonderful notebook:
|
|
https://github.com/FullStackRetrieval-com/RetrievalTutorials/blob/main/tutorials/LevelsOfTextSplitting/5_Levels_Of_Text_Splitting.ipynb
|
|
|
|
All credits to him.
|
|
|
|
At a high level, this splits into sentences, then groups into groups of 3
|
|
sentences, and then merges one that are similar in the embedding space.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
embeddings: Embeddings,
|
|
buffer_size: int = 1,
|
|
add_start_index: bool = False,
|
|
breakpoint_threshold_type: BreakpointThresholdType = "percentile",
|
|
breakpoint_threshold_amount: Optional[float] = None,
|
|
number_of_chunks: Optional[int] = None,
|
|
sentence_split_regex: str = r"(?<=[.?!])\s+",
|
|
):
|
|
self._add_start_index = add_start_index
|
|
self.embeddings = embeddings
|
|
self.buffer_size = buffer_size
|
|
self.breakpoint_threshold_type = breakpoint_threshold_type
|
|
self.number_of_chunks = number_of_chunks
|
|
self.sentence_split_regex = sentence_split_regex
|
|
if breakpoint_threshold_amount is None:
|
|
self.breakpoint_threshold_amount = BREAKPOINT_DEFAULTS[
|
|
breakpoint_threshold_type
|
|
]
|
|
else:
|
|
self.breakpoint_threshold_amount = breakpoint_threshold_amount
|
|
|
|
def _calculate_breakpoint_threshold(
|
|
self, distances: List[float]
|
|
) -> Tuple[float, List[float]]:
|
|
if self.breakpoint_threshold_type == "percentile":
|
|
return cast(
|
|
float,
|
|
np.percentile(distances, self.breakpoint_threshold_amount),
|
|
), distances
|
|
elif self.breakpoint_threshold_type == "standard_deviation":
|
|
return cast(
|
|
float,
|
|
np.mean(distances)
|
|
+ self.breakpoint_threshold_amount * np.std(distances),
|
|
), distances
|
|
elif self.breakpoint_threshold_type == "interquartile":
|
|
q1, q3 = np.percentile(distances, [25, 75])
|
|
iqr = q3 - q1
|
|
|
|
return np.mean(
|
|
distances
|
|
) + self.breakpoint_threshold_amount * iqr, distances
|
|
elif self.breakpoint_threshold_type == "gradient":
|
|
# Calculate the threshold based on the distribution of gradient of distance array. # noqa: E501
|
|
distance_gradient = np.gradient(distances, range(0, len(distances)))
|
|
return cast(
|
|
float,
|
|
np.percentile(distance_gradient, self.breakpoint_threshold_amount),
|
|
), distance_gradient
|
|
else:
|
|
raise ValueError(
|
|
f"Got unexpected `breakpoint_threshold_type`: "
|
|
f"{self.breakpoint_threshold_type}"
|
|
)
|
|
|
|
def _threshold_from_clusters(self, distances: List[float]) -> float:
|
|
"""
|
|
Calculate the threshold based on the number of chunks.
|
|
Inverse of percentile method.
|
|
"""
|
|
if self.number_of_chunks is None:
|
|
raise ValueError(
|
|
"This should never be called if `number_of_chunks` is None."
|
|
)
|
|
x1, y1 = len(distances), 0.0
|
|
x2, y2 = 1.0, 100.0
|
|
|
|
x = max(min(self.number_of_chunks, x1), x2)
|
|
|
|
# Linear interpolation formula
|
|
y = y1 + ((y2 - y1) / (x2 - x1)) * (x - x1)
|
|
y = min(max(y, 0), 100)
|
|
|
|
return cast(float, np.percentile(distances, y))
|
|
|
|
def _calculate_sentence_distances(
|
|
self, single_sentences_list: List[str]
|
|
) -> Tuple[List[float], List[dict]]:
|
|
"""Split text into multiple components."""
|
|
|
|
_sentences = [
|
|
{"sentence": x, "index": i} for i, x in enumerate(single_sentences_list)
|
|
]
|
|
sentences = combine_sentences(_sentences, self.buffer_size)
|
|
embeddings = self.embeddings.embed_documents(
|
|
[x["combined_sentence"] for x in sentences]
|
|
)
|
|
for i, sentence in enumerate(sentences):
|
|
sentence["combined_sentence_embedding"] = embeddings[i]
|
|
|
|
return calculate_cosine_distances(sentences)
|
|
|
|
def split_text(
|
|
self,
|
|
text: str,
|
|
) -> List[str]:
|
|
# Splitting the essay (by default on '.', '?', and '!')
|
|
single_sentences_list = re.split(self.sentence_split_regex, text)
|
|
|
|
# having len(single_sentences_list) == 1 would cause the following
|
|
# np.percentile to fail.
|
|
if len(single_sentences_list) == 1:
|
|
return single_sentences_list
|
|
distances, sentences = self._calculate_sentence_distances(single_sentences_list)
|
|
if self.number_of_chunks is not None:
|
|
breakpoint_distance_threshold = self._threshold_from_clusters(distances)
|
|
breakpoint_array = distances
|
|
else:
|
|
(
|
|
breakpoint_distance_threshold,
|
|
breakpoint_array,
|
|
) = self._calculate_breakpoint_threshold(distances)
|
|
|
|
indices_above_thresh = [
|
|
i
|
|
for i, x in enumerate(breakpoint_array)
|
|
if x > breakpoint_distance_threshold
|
|
]
|
|
|
|
chunks = []
|
|
start_index = 0
|
|
|
|
# Iterate through the breakpoints to slice the sentences
|
|
for index in indices_above_thresh:
|
|
# The end index is the current breakpoint
|
|
end_index = index
|
|
|
|
# Slice the sentence_dicts from the current start index to the end index
|
|
group = sentences[start_index : end_index + 1]
|
|
combined_text = " ".join([d["sentence"] for d in group])
|
|
chunks.append(combined_text)
|
|
|
|
# Update the start index for the next group
|
|
start_index = index + 1
|
|
|
|
# The last group, if any sentences remain
|
|
if start_index < len(sentences):
|
|
combined_text = " ".join([d["sentence"] for d in sentences[start_index:]])
|
|
chunks.append(combined_text)
|
|
return chunks
|
|
|
|
def create_documents(
|
|
self, texts: List[str], metadatas: Optional[List[dict]] = None
|
|
) -> List[Document]:
|
|
"""Create documents from a list of texts."""
|
|
_metadatas = metadatas or [{}] * len(texts)
|
|
documents = []
|
|
for i, text in enumerate(texts):
|
|
index = -1
|
|
for chunk in self.split_text(text):
|
|
metadata = copy.deepcopy(_metadatas[i])
|
|
if self._add_start_index:
|
|
index = text.find(chunk, index + 1)
|
|
metadata["start_index"] = index
|
|
new_doc = Document(page_content=chunk, metadata=metadata)
|
|
documents.append(new_doc)
|
|
return documents
|
|
|
|
def split_documents(self, documents: Iterable[Document]) -> List[Document]:
|
|
"""Split documents."""
|
|
texts, metadatas = [], []
|
|
for doc in documents:
|
|
texts.append(doc.page_content)
|
|
metadatas.append(doc.metadata)
|
|
return self.create_documents(texts, metadatas=metadatas)
|
|
|
|
def transform_documents(
|
|
self, documents: Sequence[Document], **kwargs: Any
|
|
) -> Sequence[Document]:
|
|
"""Transform sequence of documents by splitting them."""
|
|
return self.split_documents(list(documents))
|