omni-split 0.0.3__py3-none-any.whl

omni_split/__init__.py

@@ -0,0 +1,16 @@
+from .omni_split import OmniSplit
+from .utils.base_utils import word_preprocessing_and_return_bytesIO
+from .utils.download_test_doc import download_files_to_test_doc
+
+__version__ = "0.0.3"
+__name__ = "omni_split"
+__author__ = "dinobot22"
+
+__all__ = [
+    "__name__",
+    "__version__",
+    "__author__",
+    "OmniSplit",
+    "word_preprocessing_and_return_bytesIO",
+    "download_files_to_test_doc"
+]

omni_split/base/chonkie_base.py

@@ -0,0 +1,139 @@
+"""Base classes for chunking text."""
+
+import warnings
+from abc import ABC, abstractmethod
+from multiprocessing import Pool, cpu_count
+from typing import Any, Callable, List, Union
+
+from tqdm import tqdm
+
+from .chonkie_tokenizer import Tokenizer
+from .chonkie_types import Chunk
+
+
+class BaseChunker(ABC):
+    """Abstract base class for all chunker implementations.
+
+    All chunker implementations should inherit from this class and implement
+    the chunk() method according to their specific chunking strategy.
+    """
+
+    def __init__(self, tokenizer_or_token_counter: Union[str, Any, Callable[[str], int]]):
+        """Initialize the chunker with a tokenizer.
+
+        Args:
+            tokenizer_or_token_counter (Union[str, Any]): String, tokenizer object, or token counter object
+
+        """
+        self.tokenizer = Tokenizer(tokenizer_or_token_counter)
+
+        # Set whether to use multiprocessing or not
+        self._use_multiprocessing = True
+
+    @abstractmethod
+    def chunk(self, text: str) -> List[Chunk]:
+        """Split text into chunks according to the implementation strategy.
+
+        Args:
+            text: Input text to be chunked
+
+        Returns:
+            List of Chunk objects containing the chunked text and metadata
+
+        """
+        pass
+
+    def _determine_optimal_workers(self) -> int:
+        """Determine optimal number of workers based on system resources."""
+        try:
+            # Get CPU cores
+            cpu_cores = cpu_count()
+
+            # Never use more than 75% of available cores
+            max_workers = max(1, int(cpu_cores * 0.75))
+
+            # Cap at 8 workers
+            return min(max_workers, 8)
+
+        except Exception as e:
+            warnings.warn(f"Error determining optimal workers: {e}. Using single process.")
+            return 1
+
+    def _process_batch_sequential(self, texts: List[str], show_progress_bar: bool = True) -> List[List[Chunk]]:
+        """Process a batch of texts sequentially."""
+        return [
+            self.chunk(t)
+            for t in tqdm(
+                texts,
+                desc="🦛",
+                disable=not show_progress_bar,
+                unit="doc",
+                bar_format="{desc} ch{bar:20}nk {percentage:3.0f}% • {n_fmt}/{total_fmt} docs chunked [{elapsed}<{remaining}, {rate_fmt}] 🌱",
+                ascii=" o",
+            )
+        ]
+
+    def _process_batch_multiprocessing(self, texts: List[str], show_progress_bar: bool = True) -> List[List[Chunk]]:
+        """Process a batch of texts using multiprocessing."""
+        num_workers = self._determine_optimal_workers()
+        total = len(texts)
+        chunksize = max(1, min(total // (num_workers * 16), 10))  # Optimize chunk size
+
+        with Pool(processes=num_workers) as pool:
+            results = []
+            with tqdm(
+                total=total,
+                desc="🦛",
+                disable=not show_progress_bar,
+                unit="doc",
+                bar_format="{desc} ch{bar:20}nk {percentage:3.0f}% • {n_fmt}/{total_fmt} docs chunked [{elapsed}<{remaining}, {rate_fmt}] 🌱",
+                ascii=" o",
+            ) as pbar:
+                for result in pool.imap(self.chunk, texts, chunksize=chunksize):
+                    results.append(result)
+                    pbar.update()
+            return results
+
+    def chunk_batch(
+        self,
+        texts: List[str],
+        show_progress_bar: bool = True,
+    ) -> List[List[Chunk]]:
+        """Split a List of texts into their respective chunks.
+
+        By default, this method uses multiprocessing to parallelize the chunking process.
+
+        Args:
+            texts: List of input texts to be chunked.
+            show_progress_bar: Whether to show a progress bar.
+
+        Returns:
+            List of lists of Chunk objects containing the chunked text and metadata
+
+        """
+        if self._use_multiprocessing:
+            return self._process_batch_multiprocessing(texts, show_progress_bar)
+        else:
+            return self._process_batch_sequential(texts, show_progress_bar)
+
+    def __call__(self, text: Union[str, List[str]], show_progress_bar: bool = True) -> Union[List[Chunk], List[List[Chunk]]]:
+        """Make the chunker callable directly.
+
+        Args:
+            text: Input text or list of texts to be chunked
+            show_progress_bar: Whether to show a progress bar (for batch chunking)
+
+        Returns:
+            List of Chunk objects or list of lists of Chunk
+
+        """
+        if isinstance(text, str):
+            return self.chunk(text)
+        elif isinstance(text, list):
+            return self.chunk_batch(text, show_progress_bar)
+        else:
+            raise ValueError("Input must be a string or a list of strings.")
+
+    def __repr__(self) -> str:
+        """Return string representation of the chunker."""
+        return f"{self.__class__.__name__}()"

omni_split/base/chonkie_tokenizer.py

@@ -0,0 +1,285 @@
+"""A utility module for handling tokenization across different backends."""
+
+import importlib
+import inspect
+from collections import defaultdict
+from typing import Any, Callable, Dict, List, Union
+
+
+class Tokenizer:
+    """Unified tokenizer interface for Chonkie.
+
+    Handles tokenizer initialization and operations across different backends
+    (HuggingFace, TikToken, custom tokenizers).
+
+    Args:
+        tokenizer: Tokenizer instance or identifier (e.g., "gpt2")
+
+    Raises:
+        ImportError: If required tokenizer backend is not installed
+
+    """
+
+    def __init__(self, tokenizer: Union[str, Callable, Any] = "gpt2"):
+        """Initialize the tokenizer."""
+        # Initialize the tokenizer
+        if isinstance(tokenizer, str):
+            self.tokenizer = self._load_tokenizer(tokenizer)
+        else:
+            self.tokenizer = tokenizer
+
+        # Determine the tokenizer backend
+        self._tokenizer_backend = self._get_tokenizer_backend()
+
+    def _get_tokenizer_backend(self) -> str:
+        """Determine the tokenizer backend."""
+        # Check if the tokenizer is a character or word tokenizer
+        if "chonkie" in str(type(self.tokenizer)):
+            return "chonkie"
+        elif "transformers" in str(type(self.tokenizer)):
+            return "transformers"
+        elif "tokenizers" in str(type(self.tokenizer)):
+            return "tokenizers"
+        elif "tiktoken" in str(type(self.tokenizer)):
+            return "tiktoken"
+        elif callable(self.tokenizer) or inspect.isfunction(self.tokenizer) or inspect.ismethod(self.tokenizer):
+            return "callable"
+        else:
+            raise ValueError(f"Tokenizer backend {str(type(self.tokenizer))} not supported")
+
+    def _load_tokenizer(self, tokenizer_name: str):
+        """Load a tokenizer based on the backend."""
+        # Check if the string is equal to "character"
+        if tokenizer_name == "character":
+            return CharacterTokenizer()
+        elif tokenizer_name == "word":
+            return WordTokenizer()
+        else:
+            try:
+                if importlib.util.find_spec("tokenizers") is not None:
+                    from tokenizers import Tokenizer
+
+                    return Tokenizer.from_pretrained(tokenizer_name)
+                else:
+                    raise Warning("Tokenizers library not found. Trying tiktoken.")
+            except Exception:
+                try:
+                    if importlib.util.find_spec("tiktoken") is not None:
+                        from tiktoken import get_encoding
+
+                        return get_encoding(tokenizer_name)
+                    else:
+                        raise Warning("TikToken library not found. Trying transformers.")
+                except Exception:
+                    try:
+                        if importlib.util.find_spec("transformers") is not None:
+                            from transformers import AutoTokenizer
+
+                            return AutoTokenizer.from_pretrained(tokenizer_name)
+                        else:
+                            raise ValueError("Tokenizer not found in the following libraries: transformers, tokenizers, tiktoken")
+                    except Exception:
+                        raise ValueError("Tokenizer not found in the following libraries: transformers, tokenizers, tiktoken")
+
+    def encode(self, text: str) -> List[int]:
+        """Encode text to token ids."""
+        if self._tokenizer_backend == "chonkie":
+            return self.tokenizer.encode(text)
+        elif self._tokenizer_backend == "transformers":
+            return self.tokenizer.encode(text, add_special_tokens=False)
+        elif self._tokenizer_backend == "tokenizers":
+            return self.tokenizer.encode(text, add_special_tokens=False).ids
+        elif self._tokenizer_backend == "tiktoken":
+            return self.tokenizer.encode(text)
+        elif self._tokenizer_backend == "callable":
+            raise NotImplementedError("Callable tokenizer backend does not support encoding.")
+        else:
+            raise ValueError(f"Tokenizer backend {self._tokenizer_backend} not supported.")
+
+    def encode_batch(self, texts: List[str]) -> List[List[int]]:
+        """Encode a batch of texts."""
+        if self._tokenizer_backend == "chonkie":
+            return self.tokenizer.encode_batch(texts)
+        elif self._tokenizer_backend == "transformers":
+            return self.tokenizer.batch_encode_plus(texts, add_special_tokens=False)["input_ids"]
+        elif self._tokenizer_backend == "tokenizers":
+            return [t.ids for t in self.tokenizer.encode_batch(texts, add_special_tokens=False)]
+        elif self._tokenizer_backend == "tiktoken":
+            return self.tokenizer.encode_batch(texts)
+        elif self._tokenizer_backend == "callable":
+            raise NotImplementedError("Callable tokenizer backend does not support batch encoding.")
+        else:
+            raise ValueError(f"Tokenizer backend {self._tokenizer_backend} not supported.")
+
+    def decode(self, tokens: List[int]) -> str:
+        """Decode token ids back to text."""
+        if self._tokenizer_backend == "callable":
+            raise NotImplementedError("Callable tokenizer backend does not support decoding.")
+        return self.tokenizer.decode(tokens)
+
+    def decode_batch(self, token_lists: List[List[int]]) -> List[str]:
+        """Decode multiple token lists."""
+        if self._tokenizer_backend == "chonkie":
+            return self.tokenizer.decode_batch(token_lists)
+        elif self._tokenizer_backend == "transformers":
+            return self.tokenizer.batch_decode(token_lists, skip_special_tokens=True)
+        elif self._tokenizer_backend in ["tokenizers", "tiktoken"]:
+            return self.tokenizer.decode_batch(token_lists)
+        elif self._tokenizer_backend == "callable":
+            raise NotImplementedError("Callable tokenizer backend does not support batch decoding.")
+        else:
+            raise ValueError(f"Tokenizer backend {self._tokenizer_backend} not supported.")
+
+    def count_tokens(self, text: str) -> int:
+        """Count number of tokens in text."""
+        if self._tokenizer_backend == "chonkie":
+            return self.tokenizer.count_tokens(text)
+        elif self._tokenizer_backend == "transformers":
+            return len(self.tokenizer.encode(text, add_special_tokens=False))
+        elif self._tokenizer_backend == "tokenizers":
+            return len(self.tokenizer.encode(text, add_special_tokens=False).ids)
+        elif self._tokenizer_backend == "tiktoken":
+            return len(self.tokenizer.encode(text))
+        elif self._tokenizer_backend == "callable":
+            return self.tokenizer(text)
+        else:
+            raise ValueError(f"Tokenizer backend {self._tokenizer_backend} not supported.")
+
+    def count_tokens_batch(self, texts: List[str]) -> List[int]:
+        """Count tokens in multiple texts."""
+        if self._tokenizer_backend == "chonkie":
+            return self.tokenizer.count_tokens_batch(texts)
+        elif self._tokenizer_backend == "transformers":
+            return [len(token_list) for token_list in self.tokenizer.batch_encode_plus(texts, add_special_tokens=False)["input_ids"]]
+        elif self._tokenizer_backend == "tokenizers":
+            return [len(token_list) for token_list in [t.ids for t in self.tokenizer.encode_batch(texts, add_special_tokens=False)]]
+        elif self._tokenizer_backend == "tiktoken":
+            return [len(token_list) for token_list in self.tokenizer.encode_batch(texts)]
+        elif self._tokenizer_backend == "callable":
+            return [self.tokenizer(text) for text in texts]
+        else:
+            raise ValueError(f"Tokenizer backend {self._tokenizer_backend} not supported.")
+
+
+class CharacterTokenizer:
+    """Character-based tokenizer."""
+
+    def __init__(self):
+        """Initialize the tokenizer."""
+        # Initialize the vocabulary with a space character
+        self.vocab = []
+        self.token2id = defaultdict(lambda: len(self.vocab))
+
+        # Add space character to vocabulary
+        _ = self.token2id[" "]
+        self.vocab.append(" ")
+
+    def get_vocab(self) -> List[str]:
+        """Get the vocabulary."""
+        return self.vocab
+
+    def get_token2id(self) -> Dict[str, int]:
+        """Get the token to id mapping."""
+        return self.token2id
+
+    def encode(self, text: str) -> List[int]:
+        """Encode text to token ids."""
+        ids = []
+        for token in text:
+            token_id = self.token2id[token]
+            if token_id >= len(self.vocab):
+                self.vocab.append(token)
+            ids.append(token_id)
+        return ids
+
+    def encode_batch(self, texts: List[str]) -> List[List[int]]:
+        """Encode a batch of texts."""
+        return [self.encode(text) for text in texts]
+
+    def decode(self, tokens: List[int]) -> str:
+        """Decode token ids back to text."""
+        try:
+            return "".join([self.vocab[token] for token in tokens])
+        except IndexError:
+            raise ValueError(f"Token {tokens} not found in vocabulary.")
+
+    def decode_batch(self, token_lists: List[List[int]]) -> List[str]:
+        """Decode multiple token lists."""
+        return [self.decode(token_list) for token_list in token_lists]
+
+    def count_tokens(self, text: str) -> int:
+        """Count number of tokens in text."""
+        return len(text)
+
+    def count_tokens_batch(self, texts: List[str]) -> List[int]:
+        """Count tokens in multiple texts."""
+        return [len(text) for text in texts]
+
+    def __repr__(self) -> str:
+        """Return a string representation of the tokenizer."""
+        return f"CharacterTokenizer(vocab_size={len(self.vocab)})"
+
+
+class WordTokenizer:
+    """Word-based tokenizer."""
+
+    def __init__(self):
+        """Initialize the tokenizer."""
+        # Initialize the vocabulary with a space character
+        self.vocab = []
+        self.token2id = defaultdict(lambda: len(self.vocab))
+
+        # Add space character to vocabulary
+        _ = self.token2id[" "]
+        self.vocab.append(" ")
+
+    def get_vocab(self) -> List[str]:
+        """Get the vocabulary."""
+        return self.vocab
+
+    def get_token2id(self) -> Dict[str, int]:
+        """Get the token to id mapping."""
+        return self.token2id
+
+    def tokenize(self, text: str) -> List[str]:
+        """Tokenize text."""
+        words = text.split(" ")
+        return words
+
+    def encode(self, text: str) -> List[int]:
+        """Encode text to token ids."""
+        tokens = self.tokenize(text)
+        ids = []
+        for token in tokens:
+            token_id = self.token2id[token]
+            if token_id >= len(self.vocab):
+                self.vocab.append(token)
+            ids.append(token_id)
+        return ids
+
+    def encode_batch(self, texts: List[str]) -> List[List[int]]:
+        """Encode a batch of texts."""
+        return [self.encode(text) for text in texts]
+
+    def decode(self, tokens: List[int]) -> str:
+        """Decode token ids back to text."""
+        try:
+            return " ".join([self.vocab[token] for token in tokens])
+        except IndexError:
+            raise ValueError(f"Token {tokens} not found in vocabulary.")
+
+    def decode_batch(self, token_lists: List[List[int]]) -> List[str]:
+        """Decode multiple token lists."""
+        return [self.decode(token_list) for token_list in token_lists]
+
+    def count_tokens(self, text: str) -> int:
+        """Count number of tokens in text."""
+        return len(self.encode(text))
+
+    def count_tokens_batch(self, texts: List[str]) -> List[int]:
+        """Count tokens in multiple texts."""
+        return [len(self.encode(text)) for text in texts]
+
+    def __repr__(self) -> str:
+        """Return a string representation of the tokenizer."""
+        return f"WordTokenizer(vocab_size={len(self.vocab)})"