120 lines
4.9 KiB
Python
120 lines
4.9 KiB
Python
import logging
|
|
import uuid
|
|
import re
|
|
from typing import List, Dict, Any
|
|
from core.models import OCRPageResult, DocumentChunk
|
|
|
|
logger = logging.getLogger("MarkdownChunker")
|
|
|
|
class MarkdownChunker:
|
|
"""
|
|
Chia nhỏ văn bản (Semantic Chunking) dựa trên các thẻ Markdown (Header, Double Newline).
|
|
Theo dõi chính xác đoạn text đó thuộc trang (Page) nào.
|
|
"""
|
|
def __init__(self, max_chunk_size: int = 1500, overlap: int = 200):
|
|
self.max_chunk_size = max_chunk_size
|
|
self.overlap = overlap
|
|
|
|
def chunk_document(self,
|
|
pages: List[OCRPageResult],
|
|
metadata: Dict[str, Any]) -> List[DocumentChunk]:
|
|
"""
|
|
Nhận danh sách các trang đã được OCR/VLM dịch và trả về các Chunks.
|
|
"""
|
|
chunks = []
|
|
current_chunk_text = ""
|
|
current_page_start = 1
|
|
|
|
# 1. Ghép tất cả các trang lại kèm theo mốc (marker) trang ẩn
|
|
# Cách này giúp ta cắt văn bản liền mạch mà vẫn biết chữ nào thuộc trang nào
|
|
full_text = ""
|
|
page_markers = [] # Lưu (index_chữ, page_num)
|
|
|
|
current_char_index = 0
|
|
for page in sorted(pages, key=lambda p: p.page):
|
|
page_markers.append((current_char_index, page.page))
|
|
page_text = page.text + "\n\n"
|
|
full_text += page_text
|
|
current_char_index += len(page_text)
|
|
|
|
# 2. Cắt bằng Regex (Tách theo Markdown Heading # hoặc khoảng trắng kép \n\n)
|
|
# Tách thô các đoạn
|
|
paragraphs = re.split(r'(?=\n#{1,4}\s)', full_text) # Tách mỗi khi gặp Header
|
|
|
|
refined_paragraphs = []
|
|
for p in paragraphs:
|
|
# Nếu đoạn quá dài, cắt tiếp bằng \n\n
|
|
if len(p) > self.max_chunk_size:
|
|
sub_p = re.split(r'\n\n', p)
|
|
refined_paragraphs.extend([s.strip() for s in sub_p if s.strip()])
|
|
else:
|
|
if p.strip():
|
|
refined_paragraphs.append(p.strip())
|
|
|
|
# 3. Gộp các đoạn nhỏ thành các Chunk tối ưu
|
|
current_chunk = ""
|
|
chunk_start_index = 0
|
|
|
|
def find_page(char_index):
|
|
"""Hàm tìm số trang từ vị trí ký tự"""
|
|
found_page = page_markers[0][1]
|
|
for idx, p_num in page_markers:
|
|
if char_index >= idx:
|
|
found_page = p_num
|
|
else:
|
|
break
|
|
return found_page
|
|
|
|
char_counter = 0
|
|
for p in refined_paragraphs:
|
|
p_len = len(p)
|
|
|
|
if len(current_chunk) + p_len > self.max_chunk_size and len(current_chunk) > 0:
|
|
# Đóng gói Chunk hiện tại
|
|
p_from = find_page(chunk_start_index)
|
|
p_to = find_page(char_counter)
|
|
|
|
chunks.append(DocumentChunk(
|
|
chunk_id=f"chk_{uuid.uuid4().hex[:10]}",
|
|
file_id=metadata.get("item_id", ""),
|
|
file_name=metadata.get("name", ""),
|
|
text=current_chunk.strip(),
|
|
page_from=p_from,
|
|
page_to=p_to,
|
|
source_url=metadata.get("web_url", ""),
|
|
download_url=metadata.get("download_url", ""),
|
|
site_id=metadata.get("site_id", ""),
|
|
permissions=metadata.get("permissions", ["*"])
|
|
))
|
|
|
|
# Bắt đầu chunk mới với một chút Overlap (chối gối)
|
|
overlap_text = current_chunk[-self.overlap:] if len(current_chunk) > self.overlap else current_chunk
|
|
current_chunk = overlap_text + "\n\n" + p
|
|
chunk_start_index = char_counter - len(overlap_text)
|
|
else:
|
|
if len(current_chunk) == 0:
|
|
chunk_start_index = char_counter
|
|
current_chunk += p + "\n\n"
|
|
|
|
char_counter += p_len + 2 # +2 vì có \n\n
|
|
|
|
# Đóng gói Chunk cuối cùng
|
|
if current_chunk.strip():
|
|
p_from = find_page(chunk_start_index)
|
|
p_to = find_page(char_counter)
|
|
chunks.append(DocumentChunk(
|
|
chunk_id=f"chk_{uuid.uuid4().hex[:10]}",
|
|
file_id=metadata.get("item_id", ""),
|
|
file_name=metadata.get("name", ""),
|
|
text=current_chunk.strip(),
|
|
page_from=p_from,
|
|
page_to=p_to,
|
|
source_url=metadata.get("web_url", ""),
|
|
download_url=metadata.get("download_url", ""),
|
|
site_id=metadata.get("site_id", ""),
|
|
permissions=metadata.get("permissions", ["*"])
|
|
))
|
|
|
|
logger.info(f"Chunked document {metadata.get('name')} into {len(chunks)} chunks.")
|
|
return chunks
|