261 lines
9.8 KiB
Python
261 lines
9.8 KiB
Python
import importlib
|
|
import logging
|
|
import re
|
|
import threading
|
|
from typing import Any, Literal
|
|
|
|
from pydantic import Field
|
|
|
|
from ontocast.config import ChunkConfig
|
|
from ontocast.tool.cache import Cacher, ToolCacher
|
|
from ontocast.tool.chunk.util import SENTENCE_SPLIT_REGEX, SemanticChunker
|
|
from ontocast.tool.onto import Tool
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Optional imports for semantic chunking
|
|
torch_module: Any | None = None
|
|
embedding_model_cls: Any | None = None
|
|
try:
|
|
torch_module = importlib.import_module("torch")
|
|
langchain_huggingface_module = importlib.import_module("langchain_huggingface")
|
|
embedding_model_cls = getattr(
|
|
langchain_huggingface_module, "HuggingFaceEmbeddings", None
|
|
)
|
|
SEMANTIC_CHUNKING_AVAILABLE = embedding_model_cls is not None
|
|
except ImportError:
|
|
SEMANTIC_CHUNKING_AVAILABLE = False
|
|
|
|
|
|
class ChunkerTool(Tool):
|
|
"""Tool for semantic chunking of documents.
|
|
|
|
Falls back to naive chunking if sentence-transformers is not available.
|
|
Includes caching to avoid re-chunking the same text with the same parameters.
|
|
"""
|
|
|
|
model: str = Field(
|
|
default="sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
|
|
description="HuggingFace model name for embeddings",
|
|
)
|
|
config: ChunkConfig = Field(
|
|
default_factory=ChunkConfig, description="Chunking configuration parameters"
|
|
)
|
|
chunking_mode: Literal["semantic", "naive"] = Field(
|
|
default="semantic" if SEMANTIC_CHUNKING_AVAILABLE else "naive",
|
|
description="Chunking mode: semantic (requires sentence-transformers) or naive (fallback)",
|
|
)
|
|
cache: Any = Field(default=None, exclude=True)
|
|
|
|
def __init__(
|
|
self,
|
|
chunk_config: ChunkConfig | None = None,
|
|
cache: Cacher | None = None,
|
|
**kwargs,
|
|
):
|
|
"""Initialize the ChunkerTool.
|
|
|
|
Args:
|
|
chunk_config: Chunking configuration. If None, uses default ChunkConfig.
|
|
cache: Optional shared Cacher instance. If None, creates a new one.
|
|
**kwargs: Additional keyword arguments passed to the parent class.
|
|
"""
|
|
super().__init__(**kwargs)
|
|
self._model: Any | None = None
|
|
self._model_lock = threading.Lock() # Lock for thread-safe model initialization
|
|
|
|
# Initialize cache - use shared cacher or create new one
|
|
if cache is not None:
|
|
self.cache = ToolCacher(cache, "chunker")
|
|
else:
|
|
# Fallback for backward compatibility
|
|
shared_cache = Cacher()
|
|
self.cache = ToolCacher(shared_cache, "chunker")
|
|
|
|
# Override config if provided
|
|
if chunk_config is not None:
|
|
self.config = chunk_config
|
|
|
|
# Override chunking mode if semantic chunking is not available
|
|
if not SEMANTIC_CHUNKING_AVAILABLE and self.chunking_mode == "semantic":
|
|
self.chunking_mode = "naive"
|
|
logger.warning(
|
|
"Semantic chunking not available (sentence-transformers not installed). "
|
|
"Falling back to naive chunking."
|
|
)
|
|
|
|
def _init_model(self):
|
|
"""Initialize the embedding model in a thread-safe manner.
|
|
|
|
Uses double-checked locking pattern to ensure the model is only
|
|
initialized once, even when called concurrently from multiple threads.
|
|
"""
|
|
# Fast path: if model already initialized, return immediately
|
|
if self._model is not None:
|
|
return
|
|
|
|
# Acquire lock for thread-safe initialization
|
|
with self._model_lock:
|
|
# Double-check: another thread might have initialized it while we waited
|
|
if self._model is None and SEMANTIC_CHUNKING_AVAILABLE:
|
|
if embedding_model_cls is not None:
|
|
try:
|
|
self._model = embedding_model_cls(
|
|
model_name=self.model,
|
|
model_kwargs={
|
|
"device": "cuda"
|
|
if torch_module is not None
|
|
and torch_module.cuda.is_available()
|
|
else "cpu"
|
|
},
|
|
encode_kwargs={"normalize_embeddings": False},
|
|
)
|
|
logger.debug(f"Initialized embedding model: {self.model}")
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize embedding model: {e}")
|
|
# Set to a sentinel value to prevent repeated failed attempts
|
|
self._model = None
|
|
|
|
def _naive_chunk(self, doc: str) -> list[str]:
|
|
"""Naive chunking fallback when semantic chunking is not available.
|
|
|
|
Args:
|
|
doc: The document text to chunk.
|
|
|
|
Returns:
|
|
List of text chunks.
|
|
"""
|
|
# Split by paragraphs first (double newlines)
|
|
paragraphs = re.split(r"\n\s*\n", doc.strip())
|
|
|
|
chunks = []
|
|
current_chunk = ""
|
|
|
|
for paragraph in paragraphs:
|
|
paragraph = paragraph.strip()
|
|
if not paragraph:
|
|
continue
|
|
|
|
# If adding this paragraph would exceed max_size, start a new chunk
|
|
if (
|
|
current_chunk
|
|
and len(current_chunk) + len(paragraph) + 2 > self.config.max_size
|
|
):
|
|
if current_chunk:
|
|
chunks.append(current_chunk.strip())
|
|
current_chunk = paragraph
|
|
else:
|
|
if current_chunk:
|
|
current_chunk += "\n\n" + paragraph
|
|
else:
|
|
current_chunk = paragraph
|
|
|
|
# If a single paragraph is too large, split it by sentences
|
|
if len(current_chunk) > self.config.max_size:
|
|
# Save the previous chunk if it exists
|
|
if len(current_chunk) - len(paragraph) - 2 > 0:
|
|
prev_chunk = current_chunk[
|
|
: len(current_chunk) - len(paragraph) - 2
|
|
].strip()
|
|
if prev_chunk:
|
|
chunks.append(prev_chunk)
|
|
|
|
# Split the large paragraph by sentences
|
|
sentences = re.split(r"(?<=[.!?])\s+", paragraph)
|
|
temp_chunk = ""
|
|
|
|
for sentence in sentences:
|
|
if len(temp_chunk) + len(sentence) + 1 > self.config.max_size:
|
|
if temp_chunk:
|
|
chunks.append(temp_chunk.strip())
|
|
temp_chunk = sentence
|
|
else:
|
|
if temp_chunk:
|
|
temp_chunk += " " + sentence
|
|
else:
|
|
temp_chunk = sentence
|
|
|
|
current_chunk = temp_chunk
|
|
|
|
# Add the last chunk
|
|
if current_chunk:
|
|
chunks.append(current_chunk.strip())
|
|
|
|
# Filter out chunks that are too small
|
|
chunks = [chunk for chunk in chunks if len(chunk) >= self.config.min_size]
|
|
|
|
logger.info(f"Naive chunking produced {len(chunks)} chunks")
|
|
return chunks
|
|
|
|
def __call__(self, doc: str) -> list[str]:
|
|
"""Chunk the document using either semantic or naive chunking.
|
|
|
|
Args:
|
|
doc: The document text to chunk.
|
|
|
|
Returns:
|
|
List of text chunks.
|
|
"""
|
|
# Prepare configuration for caching
|
|
config_dict = {
|
|
"model": self.model,
|
|
"chunking_mode": self.chunking_mode,
|
|
"max_size": self.config.max_size,
|
|
"min_size": self.config.min_size,
|
|
"breakpoint_threshold_type": self.config.breakpoint_threshold_type,
|
|
"breakpoint_threshold_amount": self.config.breakpoint_threshold_amount,
|
|
}
|
|
|
|
# Check cache first
|
|
cached_result = self.cache.get(doc, config=config_dict)
|
|
if cached_result is not None:
|
|
logger.debug("Cache hit for document chunking")
|
|
return cached_result
|
|
|
|
# Perform chunking
|
|
if self.chunking_mode == "naive":
|
|
result = self._naive_chunk(doc)
|
|
else:
|
|
# Semantic chunking (requires sentence-transformers)
|
|
if not SEMANTIC_CHUNKING_AVAILABLE:
|
|
logger.warning(
|
|
"Semantic chunking requested but not available. Falling back to naive chunking."
|
|
)
|
|
result = self._naive_chunk(doc)
|
|
else:
|
|
self._init_model()
|
|
documents = [doc]
|
|
|
|
if self._model is None:
|
|
logger.warning(
|
|
"Model not initialized. Falling back to naive chunking."
|
|
)
|
|
result = self._naive_chunk(doc)
|
|
elif SemanticChunker is None:
|
|
logger.warning(
|
|
"SemanticChunker not available. Falling back to naive chunking."
|
|
)
|
|
result = self._naive_chunk(doc)
|
|
else:
|
|
text_splitter = SemanticChunker(
|
|
embeddings=self._model,
|
|
chunk_config=self.config,
|
|
sentence_split_regex=SENTENCE_SPLIT_REGEX,
|
|
)
|
|
|
|
# SemanticChunker now handles max_size internally
|
|
result_docs = text_splitter.create_documents(documents)
|
|
result = [doc.page_content for doc in result_docs]
|
|
|
|
# Log chunk lengths for debugging
|
|
lens = [len(chunk) for chunk in result]
|
|
logger.info(
|
|
f"Semantic chunking produced {len(result)} chunks with lengths: {lens}"
|
|
)
|
|
|
|
# Cache the result
|
|
self.cache.set(doc, result, config=config_dict)
|
|
logger.debug("Cached document chunking result")
|
|
|
|
return result
|