chore: import upstream snapshot with attribution
Build and push multi-arch DocsGPT Docker image / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Backend release / release (push) Has been cancelled
Bandit Security Scan / bandit_scan (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / manifest (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / manifest (push) Has been cancelled
Python linting / ruff (push) Has been cancelled
Run python tests with pytest / Run tests and count coverage (3.12) (push) Has been cancelled
React Widget Build / build (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Backend release / release (push) Has been cancelled
Bandit Security Scan / bandit_scan (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / manifest (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / manifest (push) Has been cancelled
Python linting / ruff (push) Has been cancelled
Run python tests with pytest / Run tests and count coverage (3.12) (push) Has been cancelled
React Widget Build / build (push) Has been cancelled
This commit is contained in:
@@ -0,0 +1,101 @@
|
||||
import re
|
||||
from typing import List, Tuple
|
||||
import logging
|
||||
from application.parser.chunking_creator import ChunkerCreator
|
||||
from application.parser.schema.base import Document
|
||||
from application.utils import get_encoding
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Chunker:
|
||||
"""Classic token-window chunker (registered as ``classic_chunk``).
|
||||
|
||||
Strategy dispatch lives in ``ChunkerCreator``; this class is one
|
||||
registered implementation. The ``chunking_strategy`` arg is retained for
|
||||
backward-compatible construction and is not used for dispatch here.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
chunking_strategy: str = "classic_chunk",
|
||||
max_tokens: int = 2000,
|
||||
min_tokens: int = 150,
|
||||
duplicate_headers: bool = False,
|
||||
):
|
||||
self.chunking_strategy = chunking_strategy
|
||||
self.max_tokens = max_tokens
|
||||
self.min_tokens = min_tokens
|
||||
self.duplicate_headers = duplicate_headers
|
||||
self.encoding = get_encoding()
|
||||
|
||||
def separate_header_and_body(self, text: str) -> Tuple[str, str]:
|
||||
header_pattern = r"^(.*?\n){3}"
|
||||
match = re.match(header_pattern, text)
|
||||
if match:
|
||||
header = match.group(0)
|
||||
body = text[len(header):]
|
||||
else:
|
||||
header, body = "", text # No header, treat entire text as body
|
||||
return header, body
|
||||
|
||||
|
||||
|
||||
def split_document(self, doc: Document) -> List[Document]:
|
||||
split_docs = []
|
||||
header, body = self.separate_header_and_body(doc.text)
|
||||
header_tokens = self.encoding.encode(header) if header else []
|
||||
body_tokens = self.encoding.encode(body)
|
||||
|
||||
current_position = 0
|
||||
part_index = 0
|
||||
while current_position < len(body_tokens):
|
||||
end_position = current_position + self.max_tokens - len(header_tokens)
|
||||
chunk_tokens = (header_tokens + body_tokens[current_position:end_position]
|
||||
if self.duplicate_headers or part_index == 0 else body_tokens[current_position:end_position])
|
||||
chunk_text = self.encoding.decode(chunk_tokens)
|
||||
new_doc = Document(
|
||||
text=chunk_text,
|
||||
doc_id=f"{doc.doc_id}-{part_index}",
|
||||
embedding=doc.embedding,
|
||||
extra_info={**(doc.extra_info or {}), "token_count": len(chunk_tokens)}
|
||||
)
|
||||
split_docs.append(new_doc)
|
||||
current_position = end_position
|
||||
part_index += 1
|
||||
header_tokens = []
|
||||
return split_docs
|
||||
|
||||
def classic_chunk(self, documents: List[Document]) -> List[Document]:
|
||||
processed_docs = []
|
||||
i = 0
|
||||
while i < len(documents):
|
||||
doc = documents[i]
|
||||
tokens = self.encoding.encode(doc.text)
|
||||
token_count = len(tokens)
|
||||
|
||||
if self.min_tokens <= token_count <= self.max_tokens:
|
||||
doc.extra_info = doc.extra_info or {}
|
||||
doc.extra_info["token_count"] = token_count
|
||||
processed_docs.append(doc)
|
||||
i += 1
|
||||
elif token_count < self.min_tokens:
|
||||
|
||||
doc.extra_info = doc.extra_info or {}
|
||||
doc.extra_info["token_count"] = token_count
|
||||
processed_docs.append(doc)
|
||||
i += 1
|
||||
else:
|
||||
# Split large documents
|
||||
processed_docs.extend(self.split_document(doc))
|
||||
i += 1
|
||||
return processed_docs
|
||||
|
||||
def chunk(
|
||||
self,
|
||||
documents: List[Document]
|
||||
) -> List[Document]:
|
||||
return self.classic_chunk(documents)
|
||||
|
||||
|
||||
ChunkerCreator.register("classic_chunk", Chunker)
|
||||
Reference in New Issue
Block a user