참고소스 수정본
This commit is contained in:
LASTA_DEV01\lasta
260
참고/ontocast-main/ontocast/tool/chunk/chunker.py
Normal file
260
참고/ontocast-main/ontocast/tool/chunk/chunker.py
Normal file
@@ -0,0 +1,260 @@
|
||||
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
|
||||
Reference in New Issue
Block a user