search

langroid/parsing/search.py

Utils to search for close matches in (a list of) strings. Useful for retrieval of docs/chunks relevant to a query, in the context of Retrieval-Augmented Generation (RAG), and SQLChat (e.g., to pull relevant parts of a large schema). See tests for examples: tests/main/test_string_search.py

find_fuzzy_matches_in_docs(query, docs, docs_clean, k, words_before=None, words_after=None, score_threshold=50.0)

Find approximate matches of the query in the docs and return surrounding characters.

Parameters:

Name	Type	Description	Default
query	str	The search string.	required
docs	List[Document]	List of Document objects to search through.	required
docs_clean	List[Document]	List of Document objects with cleaned content.	required
k	int	Number of best matches to return.	required
words_before	int | None	Number of words to include before each match. Default None => return max	None
words_after	int | None	Number of words to include after each match. Default None => return max	None
score_threshold	float	Keep only matches whose fuzzy-match score is STRICTLY GREATER than this value, on the 0-100 rapidfuzz scale (100 = perfect match). The default 50.0 preserves the original behavior of this function, which hard-coded the filter score > 50. A threshold that is not a finite number (e.g. None, nan, inf, -inf, or a non-numeric value) yields NO matches: the result is [].	50.0

Returns:

Type	Description
List[Tuple[Document, float]]	List[Tuple[Document,float]]: List of (Document, score) tuples.

Source code in langroid/parsing/search.py

def find_fuzzy_matches_in_docs(
    query: str,
    docs: List[Document],
    docs_clean: List[Document],
    k: int,
    words_before: int | None = None,
    words_after: int | None = None,
    score_threshold: float = 50.0,
) -> List[Tuple[Document, float]]:
    """
    Find approximate matches of the query in the docs and return surrounding
    characters.

    Args:
        query (str): The search string.
        docs (List[Document]): List of Document objects to search through.
        docs_clean (List[Document]): List of Document objects with cleaned content.
        k (int): Number of best matches to return.
        words_before (int|None): Number of words to include before each match.
            Default None => return max
        words_after (int|None): Number of words to include after each match.
            Default None => return max
        score_threshold (float): Keep only matches whose fuzzy-match score is
            STRICTLY GREATER than this value, on the 0-100 `rapidfuzz` scale
            (100 = perfect match). The default 50.0 preserves the original
            behavior of this function, which hard-coded the filter
            `score > 50`. A threshold that is not a finite number (e.g.
            `None`, `nan`, `inf`, `-inf`, or a non-numeric value) yields
            NO matches: the result is `[]`.

    Returns:
        List[Tuple[Document,float]]: List of (Document, score) tuples.
    """
    # Normalize the threshold defensively: at runtime it may arrive as
    # None or another non-numeric/non-finite value (e.g. from an
    # unvalidated config); by contract such thresholds yield no matches.
    try:
        threshold = float(score_threshold)
    except (TypeError, ValueError, OverflowError):
        return []
    if not math.isfinite(threshold):
        return []
    if len(docs) == 0:
        return []
    best_matches = process.extract(
        query,
        {i: d.content for i, d in enumerate(docs_clean)},
        limit=k,
        scorer=fuzz.partial_ratio,
    )

    # Preserve document identity when cleaned texts overlap or are identical.
    orig_doc_matches = [
        (docs[i], score) for _, score, i in best_matches if score > threshold
    ]
    if words_after is None and words_before is None:
        return orig_doc_matches
    if len(orig_doc_matches) == 0:
        return []
    if set(orig_doc_matches[0][0].model_fields) != {"content", "metadata"}:
        # If there are fields beyond just content and metadata,
        # we do NOT want to create new document objects with content fields
        # based on words_before and words_after, since we don't know how to
        # set those other fields.
        return orig_doc_matches

    contextual_matches = []
    for match, score in orig_doc_matches:
        choice_text = match.content
        contexts = []
        while choice_text != "":
            context, start_pos, end_pos = get_context(
                query, choice_text, words_before, words_after
            )
            if context == "" or end_pos == 0:
                break
            contexts.append(context)
            words = choice_text.split()
            end_pos = min(end_pos, len(words))
            choice_text = " ".join(words[end_pos:])
        if len(contexts) > 0:
            contextual_matches.append(
                (
                    Document(
                        content=" ... ".join(contexts),
                        metadata=match.metadata,
                    ),
                    score,
                )
            )

    return contextual_matches

preprocess_text(text)

Preprocesses the given text by: 1. Lowercasing all words. 2. Tokenizing (splitting the text into words). 3. Removing punctuation. 4. Removing stopwords. 5. Lemmatizing words.

Parameters:

Name	Type	Description	Default
text	str	The input text.	required

Returns:

Name	Type	Description
str	str	The preprocessed text.

Source code in langroid/parsing/search.py

def preprocess_text(text: str) -> str:
    """
    Preprocesses the given text by:
    1. Lowercasing all words.
    2. Tokenizing (splitting the text into words).
    3. Removing punctuation.
    4. Removing stopwords.
    5. Lemmatizing words.

    Args:
        text (str): The input text.

    Returns:
        str: The preprocessed text.
    """
    # Ensure the NLTK resources are available
    for resource in ["tokenizers/punkt", "corpora/wordnet", "corpora/stopwords"]:
        download_nltk_resource(resource)
    from nltk.corpus import stopwords
    from nltk.stem import WordNetLemmatizer
    from nltk.tokenize import RegexpTokenizer

    # Lowercase the text
    text = text.lower()

    # Tokenize the text and remove punctuation
    tokenizer = RegexpTokenizer(r"\w+")
    tokens = tokenizer.tokenize(text)

    # Remove stopwords
    stop_words = set(stopwords.words("english"))
    tokens = [t for t in tokens if t not in stop_words]

    # Lemmatize words
    lemmatizer = WordNetLemmatizer()
    tokens = [lemmatizer.lemmatize(t) for t in tokens]

    # Join the words back into a string
    text = " ".join(tokens)

    return text

find_closest_matches_with_bm25(docs, docs_clean, query, k=5)

Finds the k closest approximate matches using the BM25 algorithm.

Parameters:

Name	Type	Description	Default
docs	List[Document]	List of Documents to search through.	required
docs_clean	List[Document]	List of cleaned Documents	required
query	str	The search query.	required
k	int	Number of matches to retrieve. Defaults to 5.	5

Returns:

Type	Description
List[Tuple[Document, float]]	List[Tuple[Document,float]]: List of (Document, score) tuples.

Source code in langroid/parsing/search.py

def find_closest_matches_with_bm25(
    docs: List[Document],
    docs_clean: List[Document],
    query: str,
    k: int = 5,
) -> List[Tuple[Document, float]]:
    """
    Finds the k closest approximate matches using the BM25 algorithm.

    Args:
        docs (List[Document]): List of Documents to search through.
        docs_clean (List[Document]): List of cleaned Documents
        query (str): The search query.
        k (int, optional): Number of matches to retrieve. Defaults to 5.

    Returns:
        List[Tuple[Document,float]]: List of (Document, score) tuples.
    """
    if len(docs) == 0:
        return []
    texts = [doc.content for doc in docs_clean]
    query = preprocess_text(query)

    text_words = [text.split() for text in texts]

    bm25 = BM25Okapi(text_words)
    query_words = query.split()
    doc_scores = bm25.get_scores(query_words)

    # Get indices of top k scores
    top_indices = sorted(range(len(doc_scores)), key=lambda i: -doc_scores[i])[:k]

    # return the original docs, based on the scores from cleaned docs
    return [(docs[i], doc_scores[i]) for i in top_indices]

get_context(query, text, words_before=100, words_after=100)

Returns a portion of text containing the best approximate match of the query, including b words before and a words after the match.

Args: query (str): The string to search for. text (str): The body of text in which to search. words_before (int|None): Number of words before the query to return. None includes all preceding words; 0 adds no preceding words. words_after (int|None): Number of words after the query to return. None includes all following words; 0 adds no following words.

str: A string containing b words before, the match, and a words after the best approximate match position of the query in the text. The text is extracted from the original text, preserving formatting, whitespace, etc, so it does not disturb any downstream processing. If no match is found, returns empty string. int: The start position of the match in the text. int: The end position of the match in the text.

Example:

get_context("apple", "The quick brown fox jumps over the apple.", 3, 2)

'fox jumps over the apple.'

Source code in langroid/parsing/search.py

def get_context(
    query: str,
    text: str,
    words_before: int | None = 100,
    words_after: int | None = 100,
) -> Tuple[str, int, int]:
    """
    Returns a portion of text containing the best approximate match of the query,
    including b words before and a words after the match.

    Args:
    query (str): The string to search for.
    text (str): The body of text in which to search.
    words_before (int|None): Number of words before the query to return.
        None includes all preceding words; 0 adds no preceding words.
    words_after (int|None): Number of words after the query to return.
        None includes all following words; 0 adds no following words.

    Returns:
    str: A string containing b words before, the match, and a words after
        the best approximate match position of the query in the text.
        The text is extracted from the original `text`, preserving formatting,
        whitespace, etc, so it does not disturb any downstream processing.
        If no match is found, returns empty string.
    int: The start position of the match in the text.
    int: The end position of the match in the text.

    Example:
    >>> get_context("apple", "The quick brown fox jumps over the apple.", 3, 2)
    # 'fox jumps over the apple.'
    """

    # If no word limits specified, return full text
    if words_after is None and words_before is None:
        # return entire text since we're not asked to return a bounded context
        return text, 0, 0

    # make sure there is a good enough match to the query
    if fuzz.partial_ratio(query, text) < 40:
        return "", 0, 0

    # Find best matching position of query in text
    sequence_matcher = difflib.SequenceMatcher(None, text, query)
    match = sequence_matcher.find_longest_match(0, len(text), 0, len(query))

    if match.size == 0:
        return "", 0, 0

    # Find character positions where words start
    segments = text.split()
    n_segs = len(segments)
    word_positions = [m.start() for m in re.finditer(r"\S+", text)]

    # Locate the word CONTAINING the match start; if the match starts in
    # whitespace, use the word that follows it. Counting the words in
    # text[:match.a] would be off by one whenever the match begins mid-word
    # (e.g. after a leading "(" or at a case difference), since the partial
    # prefix is counted as a whole word.
    start_segment_pos = bisect.bisect_right(word_positions, match.a) - 1
    if start_segment_pos < 0 or match.a >= (
        word_positions[start_segment_pos] + len(segments[start_segment_pos])
    ):
        # match.a lies in whitespace: advance to the next word
        start_segment_pos = min(start_segment_pos + 1, n_segs - 1)

    # Calculate word window boundaries
    words_before = n_segs if words_before is None else words_before
    words_after = n_segs if words_after is None else words_after
    start_pos = max(0, start_segment_pos - words_before)
    end_pos = min(len(segments), start_segment_pos + words_after + len(query.split()))

    # Convert word positions to character positions
    start_char = word_positions[start_pos] if start_pos < len(word_positions) else 0
    end_char = word_positions[end_pos - 1] + len(segments[end_pos - 1])

    # return exact substring with original formatting
    return text[start_char:end_char], start_pos, end_pos

eliminate_near_duplicates(passages, threshold=0.8)

Eliminate near duplicate text passages from a given list using MinHash and LSH. TODO: this has not been tested and the datasketch lib is not a dependency. Args: passages (List[str]): A list of text passages. threshold (float, optional): Jaccard similarity threshold to consider two passages as near-duplicates. Default is 0.8.

Returns:

Type	Description
List[str]	List[str]: A list of passages after eliminating near duplicates.

Example

passages = ["Hello world", "Hello, world!", "Hi there", "Hello world!"] print(eliminate_near_duplicates(passages))

['Hello world', 'Hi there']

Source code in langroid/parsing/search.py

def eliminate_near_duplicates(passages: List[str], threshold: float = 0.8) -> List[str]:
    """
    Eliminate near duplicate text passages from a given list using MinHash and LSH.
    TODO: this has not been tested and the datasketch lib is not a dependency.
    Args:
        passages (List[str]): A list of text passages.
        threshold (float, optional): Jaccard similarity threshold to consider two
                                     passages as near-duplicates. Default is 0.8.

    Returns:
        List[str]: A list of passages after eliminating near duplicates.

    Example:
        passages = ["Hello world", "Hello, world!", "Hi there", "Hello world!"]
        print(eliminate_near_duplicates(passages))
        # ['Hello world', 'Hi there']
    """

    from datasketch import MinHash, MinHashLSH

    # Create LSH index
    lsh = MinHashLSH(threshold=threshold, num_perm=128)

    # Create MinHash objects for each passage and insert to LSH
    minhashes = {}
    for idx, passage in enumerate(passages):
        m = MinHash(num_perm=128)
        for word in passage.split():
            m.update(word.encode("utf-8"))
        lsh.insert(idx, m)
        minhashes[idx] = m

    unique_idxs = set()
    for idx in minhashes.keys():
        # Query for similar passages (including itself)
        result = lsh.query(minhashes[idx])

        # If only the passage itself is returned, it's unique
        if len(result) == 1 and idx in result:
            unique_idxs.add(idx)

    return [passages[idx] for idx in unique_idxs]