codekingpro/portable-devtools

1"""Spacy text splitter."""2 3from __future__ import annotations4 5from typing import TYPE_CHECKING, Any6 7from typing_extensions import override8 9from langchain_text_splitters.base import TextSplitter10 11try:12    # Type ignores needed as long as spacy doesn't support Python 3.14.13    import spacy  # type: ignore[import-not-found, unused-ignore]14    from spacy.lang.en import English  # type: ignore[import-not-found, unused-ignore]15 16    if TYPE_CHECKING:17        from spacy.language import (  # type: ignore[import-not-found, unused-ignore]18            Language,19        )20 21    _HAS_SPACY = True22except ImportError:23    _HAS_SPACY = False24 25 26class SpacyTextSplitter(TextSplitter):27    """Splitting text using Spacy package.28 29    Per default, Spacy's `en_core_web_sm` model is used and30    its default max_length is 1000000 (it is the length of maximum character31    this model takes which can be increased for large files). For a faster, but32    potentially less accurate splitting, you can use `pipeline='sentencizer'`.33    """34 35    def __init__(36        self,37        separator: str = "\n\n",38        pipeline: str = "en_core_web_sm",39        max_length: int = 1_000_000,40        *,41        strip_whitespace: bool = True,42        **kwargs: Any,43    ) -> None:44        """Initialize the spacy text splitter."""45        super().__init__(**kwargs)46        self._tokenizer = _make_spacy_pipeline_for_splitting(47            pipeline, max_length=max_length48        )49        self._separator = separator50        self._strip_whitespace = strip_whitespace51 52    @override53    def split_text(self, text: str) -> list[str]:54        splits = (55            s.text if self._strip_whitespace else s.text_with_ws56            for s in self._tokenizer(text).sents57        )58        return self._merge_splits(splits, self._separator)59 60 61def _make_spacy_pipeline_for_splitting(62    pipeline: str, *, max_length: int = 1_000_00063) -> Language:64    if not _HAS_SPACY:65        msg = "Spacy is not installed, please install it with `pip install spacy`."66        raise ImportError(msg)67    if pipeline == "sentencizer":68        sentencizer: Language = English()69        sentencizer.add_pipe("sentencizer")70    else:71        sentencizer = spacy.load(pipeline, exclude=["ner", "tagger"])72        sentencizer.max_length = max_length73    return sentencizer74