Spaces:

aliroya32
/

ContiAI

Sleeping

ContiAI / rag /chuncking.py

Add login page

325b94c 7 months ago

2.78 kB

# chunking.py
import re
from typing import List, Tuple


def approx_token_count(text: str) -> int:
    """
    Rough token estimate for chunk sizing (heuristic).
    Keep this simple but consistent: ~4 chars per token.
    """
    return max(1, len(text) // 4)


def chunk_pages(
    pages: List[str],
    target_tokens: int = 520,
    overlap_tokens: int = 80,
) -> List[Tuple[str, int, int]]:
    """
    Split pages (list[str]) into overlapping chunks.
    Returns list of tuples: (chunk_text, page_start, page_end)
    Overlap is implemented at paragraph level (keeps page ranges correct).
    """
    chunks: List[Tuple[str, int, int]] = []

    buffer: List[str] = []
    buffer_pages: List[int] = []
    buffer_tokens = 0

    def make_chunk():
        nonlocal buffer, buffer_pages, buffer_tokens
        if not buffer:
            return
        chunk_text = "\n\n".join(buffer).strip()
        page_start = min(buffer_pages)
        page_end = max(buffer_pages)
        chunks.append((chunk_text, page_start, page_end))

    for page_idx, page in enumerate(pages, start=1):
        paragraphs = [p.strip() for p in re.split(r"\n\s*\n", page) if p.strip()]

        for para in paragraphs:
            t = approx_token_count(para)

            # If single paragraph exceeds target, create it as its own chunk
            if buffer_tokens == 0 and t > target_tokens:
                # make chunk with this large paragraph alone
                buffer = [para]
                buffer_pages = [page_idx]
                buffer_tokens = t
                make_chunk()
                buffer, buffer_pages, buffer_tokens = [], [], 0
                continue

            # If adding this paragraph would exceed target, flush current chunk
            if buffer_tokens + t > target_tokens:
                make_chunk()

                # prepare overlap: keep tail paragraphs whose tokens sum >= overlap_tokens
                tail_buffer: List[str] = []
                tail_pages: List[int] = []
                tail_tokens = 0
                # iterate buffer in reverse to pick tail paragraphs
                for p, p_pg in zip(reversed(buffer), reversed(buffer_pages)):
                    pt = approx_token_count(p)
                    tail_buffer.insert(0, p)
                    tail_pages.insert(0, p_pg)
                    tail_tokens += pt
                    if tail_tokens >= overlap_tokens:
                        break

                buffer = tail_buffer
                buffer_pages = tail_pages
                buffer_tokens = tail_tokens

            # append current paragraph
            buffer.append(para)
            buffer_pages.append(page_idx)
            buffer_tokens += t

    # final flush
    if buffer:
        make_chunk()

    return chunks

	# chunking.py
	import re
	from typing import List, Tuple


	def approx_token_count(text: str) -> int:
	"""
	Rough token estimate for chunk sizing (heuristic).
	Keep this simple but consistent: ~4 chars per token.
	"""
	return max(1, len(text) // 4)


	def chunk_pages(
	pages: List[str],
	target_tokens: int = 520,
	overlap_tokens: int = 80,
	) -> List[Tuple[str, int, int]]:
	"""
	Split pages (list[str]) into overlapping chunks.
	Returns list of tuples: (chunk_text, page_start, page_end)

	Overlap is implemented at paragraph level (keeps page ranges correct).
	"""
	chunks: List[Tuple[str, int, int]] = []

	buffer: List[str] = []
	buffer_pages: List[int] = []
	buffer_tokens = 0

	def make_chunk():
	nonlocal buffer, buffer_pages, buffer_tokens
	if not buffer:
	return
	chunk_text = "\n\n".join(buffer).strip()
	page_start = min(buffer_pages)
	page_end = max(buffer_pages)
	chunks.append((chunk_text, page_start, page_end))

	for page_idx, page in enumerate(pages, start=1):
	paragraphs = [p.strip() for p in re.split(r"\n\s*\n", page) if p.strip()]

	for para in paragraphs:
	t = approx_token_count(para)

	# If single paragraph exceeds target, create it as its own chunk
	if buffer_tokens == 0 and t > target_tokens:
	# make chunk with this large paragraph alone
	buffer = [para]
	buffer_pages = [page_idx]
	buffer_tokens = t
	make_chunk()
	buffer, buffer_pages, buffer_tokens = [], [], 0
	continue

	# If adding this paragraph would exceed target, flush current chunk
	if buffer_tokens + t > target_tokens:
	make_chunk()

	# prepare overlap: keep tail paragraphs whose tokens sum >= overlap_tokens
	tail_buffer: List[str] = []
	tail_pages: List[int] = []
	tail_tokens = 0
	# iterate buffer in reverse to pick tail paragraphs
	for p, p_pg in zip(reversed(buffer), reversed(buffer_pages)):
	pt = approx_token_count(p)
	tail_buffer.insert(0, p)
	tail_pages.insert(0, p_pg)
	tail_tokens += pt
	if tail_tokens >= overlap_tokens:
	break

	buffer = tail_buffer
	buffer_pages = tail_pages
	buffer_tokens = tail_tokens

	# append current paragraph
	buffer.append(para)
	buffer_pages.append(page_idx)
	buffer_tokens += t

	# final flush
	if buffer:
	make_chunk()

	return chunks