sec2md 0.1.5__py3-none-any.whl

sec2md/chunker/markdown_blocks.py

@@ -0,0 +1,135 @@
+import re
+from typing import List
+from pydantic import BaseModel, Field, computed_field
+
+try:
+    import tiktoken
+    TIKTOKEN_AVAILABLE = True
+except ImportError:
+    TIKTOKEN_AVAILABLE = False
+
+
+def estimate_tokens(text: str) -> int:
+    """
+    Calculate token count for text.
+
+    Uses tiktoken with cl100k_base encoding (gpt-3.5-turbo/gpt-4) if available.
+    Falls back to character/4 heuristic if tiktoken is not installed.
+    """
+    if TIKTOKEN_AVAILABLE:
+        encoding = tiktoken.get_encoding("cl100k_base")
+        return len(encoding.encode(text))
+    else:
+        # Fallback: simple heuristic
+        return max(1, len(text) // 4)
+
+
+def split_sentences(text: str) -> List[str]:
+    """Simple regex-based sentence splitter"""
+    # Split on .!? followed by whitespace and capital letter or end of string
+    # Handles common abbreviations like Mr., Dr., Inc., etc.
+    sentences = re.split(r'(?<=[.!?])\s+(?=[A-Z])', text)
+    return [s.strip() for s in sentences if s.strip()]
+
+
+class BaseBlock(BaseModel):
+    """Base class for markdown blocks."""
+    block_type: str = Field(..., description="Type of markdown block")
+    content: str = Field(..., description="Block content")
+    page: int = Field(..., description="Page number")
+
+    model_config = {"frozen": False}
+
+    @computed_field
+    @property
+    def tokens(self) -> int:
+        return estimate_tokens(self.content)
+
+
+class Sentence(BaseModel):
+    """Sentence within a text block."""
+    content: str = Field(..., description="Sentence content")
+
+    model_config = {"frozen": False}
+
+    @computed_field
+    @property
+    def tokens(self) -> int:
+        return estimate_tokens(self.content)
+
+
+class TextBlock(BaseBlock):
+    block_type: str = Field(default='Text', description="Text block type")
+
+    @computed_field
+    @property
+    def sentences(self) -> List[Sentence]:
+        """Returns the text block sentences"""
+        return [Sentence(content=content) for content in split_sentences(self.content)]
+
+    @classmethod
+    def from_sentences(cls, sentences: List[Sentence], page: int):
+        content = " ".join([sentence.content for sentence in sentences])
+        return cls(content=content, page=page, block_type='Text')
+
+
+class AudioParagraphBlock(BaseBlock):
+    block_type: str = Field(default="Text", description="Audio paragraph block type")
+    paragraph_id: int = Field(..., description="Paragraph ID")
+    audio_start: float = Field(..., description="Audio start time")
+    audio_end: float = Field(..., description="Audio end time")
+
+    @computed_field
+    @property
+    def sentences(self) -> List[Sentence]:
+        """Returns the text block sentences"""
+        return [Sentence(content=content) for content in split_sentences(self.content)]
+
+    def format(self) -> dict:
+        """Formats the audio paragraphs"""
+        return {"id": self.paragraph_id, "content": self.content, "start": self.audio_start, "end": self.audio_end}
+
+
+class TableBlock(BaseModel):
+    block_type: str = Field(default='Table', description="Table block type")
+    content: str = Field(..., description="Table content")
+    page: int = Field(..., description="Page number")
+
+    model_config = {"frozen": False}
+
+    def __init__(self, **data):
+        if 'content' in data:
+            data['content'] = self._to_minified_markdown_static(data['content'])
+        super().__init__(**data)
+
+    @computed_field
+    @property
+    def tokens(self) -> int:
+        return estimate_tokens(self.content)
+
+    @staticmethod
+    def _to_minified_markdown_static(content: str) -> str:
+        """Returns the table in a Minified Markdown format"""
+        lines = content.split('\n')
+        cleaned_lines = []
+
+        for i, line in enumerate(lines):
+            if not line.strip():
+                continue
+
+            parts = line.split('|')
+            cleaned_parts = [re.sub(r'\s+', ' ', part.strip()) for part in parts]
+            cleaned_line = '|'.join(cleaned_parts)
+
+            if i == 1:
+                num_cols = len(cleaned_parts) - 1
+                separator = '|' + '|'.join(['---'] * num_cols) + '|'
+                cleaned_lines.append(separator)
+            else:
+                cleaned_lines.append(cleaned_line)
+
+        return '\n'.join(cleaned_lines)
+
+
+class HeaderBlock(BaseBlock):
+    block_type: str = Field(default='Header', description="Header block type")

sec2md/chunker/markdown_chunk.py

@@ -0,0 +1,133 @@
+from typing import List, Optional, Tuple, TYPE_CHECKING
+from pydantic import BaseModel, Field, computed_field
+
+from sec2md.chunker.markdown_blocks import BaseBlock
+
+if TYPE_CHECKING:
+    from sec2md.models import Element
+else:
+    Element = 'Element'  # Forward reference for Pydantic
+
+
+class MarkdownChunk(BaseModel):
+    """Represents a chunk of markdown content that can be embedded"""
+
+    blocks: List[BaseBlock] = Field(..., description="List of markdown blocks in this chunk")
+    header: Optional[str] = Field(None, description="Optional header for embedding context")
+    elements: List['Element'] = Field(default_factory=list, description="Element objects for citation")
+    vector: Optional[List[float]] = Field(None, description="Vector embedding for this chunk")
+
+    model_config = {"frozen": False, "arbitrary_types_allowed": True}
+
+    @computed_field
+    @property
+    def page(self) -> int:
+        """First page (for backward compatibility)."""
+        return self.blocks[0].page if self.blocks else 1
+
+    def set_vector(self, vector: List[float]):
+        """Set the vector embedding for this chunk"""
+        self.vector = vector
+
+    @computed_field
+    @property
+    def start_page(self) -> int:
+        """First page this chunk appears on (from elements or blocks)."""
+        # Prefer elements since they have actual page info from the document
+        if self.elements:
+            return min(e.page_start for e in self.elements)
+        elif self.blocks:
+            return min(block.page for block in self.blocks)
+        return self.page
+
+    @computed_field
+    @property
+    def end_page(self) -> int:
+        """Last page this chunk appears on (from elements or blocks)."""
+        # Prefer elements since they have actual page info from the document
+        if self.elements:
+            return max(e.page_end for e in self.elements)
+        elif self.blocks:
+            return max(block.page for block in self.blocks)
+        return self.page
+
+    @computed_field
+    @property
+    def page_range(self) -> Tuple[int, int]:
+        """(start_page, end_page) tuple."""
+        return (self.start_page, self.end_page)
+
+    @computed_field
+    @property
+    def content(self) -> str:
+        """Get the text content of this chunk"""
+        return "\n".join([block.content for block in self.blocks])
+
+    @computed_field
+    @property
+    def data(self) -> List[dict]:
+        """Returns a list of block data grouped by page with ONLY the chunk's content"""
+        page_blocks = {}
+
+        for block in self.blocks:
+            if block.page not in page_blocks:
+                page_blocks[block.page] = []
+            page_blocks[block.page].append(block)
+
+        page_content_data = []
+        for page, blocks in page_blocks.items():
+            # Only include the content from blocks in THIS chunk, not full page content
+            page_content = "\n".join(block.content for block in blocks)
+            if not page_content.strip():
+                continue
+
+            page_content_data.append({
+                "page": page,
+                "content": page_content
+            })
+
+        return sorted(page_content_data, key=lambda x: x["page"])
+
+    @computed_field
+    @property
+    def pages(self) -> List[dict]:
+        """Returns a list of pages with ONLY this chunk's content (not full page content)"""
+        return self.data
+
+    @computed_field
+    @property
+    def embedding_text(self) -> str:
+        """Get the text to use for embedding, with optional header prepended"""
+        if self.header:
+            return f"{self.header}\n\n...\n\n{self.content}"
+        return self.content
+
+    @computed_field
+    @property
+    def has_table(self) -> bool:
+        """Returns True if this chunk contains one or more table blocks"""
+        return any(block.block_type == 'Table' for block in self.blocks)
+
+    @computed_field
+    @property
+    def num_tokens(self) -> int:
+        """Returns the total number of tokens in this chunk"""
+        return sum(block.tokens for block in self.blocks)
+
+    @computed_field
+    @property
+    def element_ids(self) -> List[str]:
+        """List of element IDs for citations."""
+        return [e.id for e in self.elements] if self.elements else []
+
+    def to_dict(self) -> dict:
+        """Alias for model_dump() - kept for backward compat during alpha."""
+        return self.model_dump()
+
+    def __repr__(self):
+        pages_str = f"{self.start_page}-{self.end_page}" if self.start_page != self.end_page else str(self.start_page)
+        return f"MarkdownChunk(pages={pages_str}, blocks={len(self.blocks)}, tokens={self.num_tokens})"
+
+    def _repr_markdown_(self):
+        """This method is called by IPython to display as Markdown"""
+        return self.content

sec2md/chunker/markdown_chunker.py

@@ -0,0 +1,270 @@
+import logging
+from typing import Union, Tuple, List, Dict, Any
+
+from sec2md.chunker.markdown_chunk import MarkdownChunk
+from sec2md.chunker.markdown_blocks import BaseBlock, TextBlock, TableBlock, HeaderBlock
+
+# Rebuild MarkdownChunk after Element is defined
+from sec2md.models import Element
+MarkdownChunk.model_rebuild()
+
+logger = logging.getLogger(__name__)
+
+
+class MarkdownChunker:
+    """Splits markdown content into chunks"""
+
+    def __init__(self, chunk_size: int = 512, chunk_overlap: int = 128):
+        self.chunk_size = chunk_size
+        self.chunk_overlap = chunk_overlap
+
+    def split(self, pages: List[Any], header: str = None) -> List[MarkdownChunk]:
+        """Split the pages into chunks with optional header for embedding context.
+
+        Args:
+            pages: List of Page objects
+            header: Optional header to prepend to each chunk's embedding_text
+
+        Returns:
+            List of MarkdownChunk objects
+        """
+        # Build element map: page -> List[Element objects]
+        page_elements = {}
+        for page in pages:
+            if hasattr(page, 'elements') and page.elements:
+                page_elements[page.number] = page.elements
+
+        blocks = self._split_into_blocks(pages=pages)
+        return self._chunk_blocks(blocks=blocks, header=header, page_elements=page_elements)
+
+    def chunk_text(self, text: str) -> List[str]:
+        """Chunk a single text string into multiple chunks"""
+        from sec2md.models import Page
+        pages = [Page(number=0, content=text)]
+        chunks = self.split(pages=pages)
+        return [chunk.content for chunk in chunks]
+
+    @staticmethod
+    def _split_into_blocks(pages: List[Any]):
+        """Splits the pages into blocks.
+
+        Args:
+            pages: List of Page objects
+
+        Returns:
+            List of BaseBlock objects
+        """
+        from sec2md.models import Page
+
+        blocks = []
+        table_content = ""
+        last_page = None
+
+        for page in pages:
+            last_page = page
+
+            for line in page.content.split('\n'):
+                if table_content and not MarkdownChunker._is_table_line(line):
+                    blocks.append(TableBlock(content=table_content, page=page.number))
+                    table_content = ""
+
+                if line.startswith("#"):
+                    blocks.append(HeaderBlock(content=line, page=page.number))
+
+                elif MarkdownChunker._is_table_line(line):
+                    table_content += f"{line}\n"
+
+                else:
+                    blocks.append(TextBlock(content=line, page=page.number))
+
+        if table_content and last_page:
+            blocks.append(TableBlock(content=table_content, page=last_page.number))
+
+        return blocks
+
+    @staticmethod
+    def _is_table_line(line: str) -> bool:
+        import re
+        if '|' not in line:
+            return False
+        stripped = line.strip()
+        if not stripped:
+            return False
+        align_pattern = re.compile(r'^\s*:?-+:?\s*$')
+        cells = [c.strip() for c in stripped.strip('|').split('|')]
+        if all(align_pattern.match(c) for c in cells):
+            return True
+        return True
+
+    def _chunk_blocks(self, blocks: List[BaseBlock], header: str = None, page_elements: dict = None) -> List[MarkdownChunk]:
+        """Converts the blocks to chunks"""
+        page_elements = page_elements or {}
+        chunks = []
+        chunk_blocks = []
+        num_tokens = 0
+
+        for i, block in enumerate(blocks):
+            next_block = blocks[i + 1] if i + 1 < len(blocks) else None
+
+            if block.block_type == 'Text':
+                chunk_blocks, num_tokens, chunks = self._process_text_block(
+                    block, chunk_blocks, num_tokens, chunks, header, page_elements
+                )
+
+            elif block.block_type == 'Table':
+                chunk_blocks, num_tokens, chunks = self._process_table_block(
+                    block, chunk_blocks, num_tokens, chunks, blocks, i, header, page_elements
+                )
+
+            else:
+                chunk_blocks, num_tokens, chunks = self._process_header_table_block(
+                    block, chunk_blocks, num_tokens, chunks, next_block, header, page_elements
+                )
+
+        if chunk_blocks:
+            self._finalize_chunk(chunks, chunk_blocks, header, page_elements)
+
+        return chunks
+
+    def _process_text_block(self, block: TextBlock, chunk_blocks: List[BaseBlock], num_tokens: int,
+                            chunks: List[MarkdownChunk], header: str = None, page_elements: dict = None):
+        """Process a text block by breaking it into sentences if needed"""
+        sentences = []
+        sentences_tokens = 0
+
+        for sentence in block.sentences:
+            if num_tokens + sentences_tokens + sentence.tokens > self.chunk_size:
+                if sentences:
+                    new_block = TextBlock.from_sentences(sentences=sentences, page=block.page)
+                    chunk_blocks.append(new_block)
+                    num_tokens += sentences_tokens
+
+                chunks, chunk_blocks, num_tokens = self._create_chunk(chunks=chunks, blocks=chunk_blocks, header=header, page_elements=page_elements)
+
+                sentences = [sentence]
+                sentences_tokens = sentence.tokens
+
+            else:
+                sentences.append(sentence)
+                sentences_tokens += sentence.tokens
+
+        if sentences:
+            new_block = TextBlock.from_sentences(sentences=sentences, page=block.page)
+            chunk_blocks.append(new_block)
+            num_tokens += sentences_tokens
+
+        return chunk_blocks, num_tokens, chunks
+
+    def _process_table_block(self, block: BaseBlock, chunk_blocks: List[BaseBlock], num_tokens: int,
+                             chunks: List[MarkdownChunk], all_blocks: List[BaseBlock], block_idx: int, header: str = None, page_elements: dict = None):
+        """Process a table block with optional header backtrack"""
+        context = []
+        context_tokens = 0
+
+        # Backtrack for header only if 1-2 short blocks precede
+        count = 0
+        for j in range(block_idx - 1, -1, -1):
+            prev = all_blocks[j]
+            if prev.page != block.page:
+                break
+            if prev.block_type == 'Header':
+                if context_tokens + prev.tokens <= 128:
+                    context.insert(0, prev)
+                    context_tokens += prev.tokens
+                break
+            elif prev.block_type == 'Text' and prev.content.strip():
+                count += 1
+                if count > 2:
+                    break
+                if context_tokens + prev.tokens <= 128:
+                    context.insert(0, prev)
+                    context_tokens += prev.tokens
+                else:
+                    break
+
+        if num_tokens + context_tokens + block.tokens > self.chunk_size:
+            if chunk_blocks:
+                chunks, chunk_blocks, num_tokens = self._create_chunk(chunks=chunks, blocks=chunk_blocks, header=header, page_elements=page_elements)
+
+            # If we're backtracking context and the last chunk is ONLY that context, remove it
+            if context and chunks and len(chunks[-1].blocks) == len(context):
+                if all(chunks[-1].blocks[i] == context[i] for i in range(len(context))):
+                    chunks.pop()
+
+            chunk_blocks = context + [block]
+            num_tokens = context_tokens + block.tokens
+        else:
+            chunk_blocks.extend(context + [block])
+            num_tokens += context_tokens + block.tokens
+
+        return chunk_blocks, num_tokens, chunks
+
+    def _process_header_table_block(self, block: BaseBlock, chunk_blocks: List[BaseBlock], num_tokens: int,
+                                    chunks: List[MarkdownChunk], next_block: BaseBlock, header: str = None, page_elements: dict = None):
+        """Process a header block"""
+        if not chunk_blocks:
+            chunk_blocks.append(block)
+            num_tokens += block.tokens
+            return chunk_blocks, num_tokens, chunks
+
+        # Don't split if current content is small and next is a table
+        if next_block and next_block.block_type == 'Table' and num_tokens < self.chunk_overlap:
+            chunk_blocks.append(block)
+            num_tokens += block.tokens
+            return chunk_blocks, num_tokens, chunks
+
+        if num_tokens + block.tokens > self.chunk_size:
+            chunks, chunk_blocks, num_tokens = self._create_chunk(chunks=chunks, blocks=chunk_blocks, header=header, page_elements=page_elements)
+            chunk_blocks.append(block)
+            num_tokens += block.tokens
+        else:
+            chunk_blocks.append(block)
+            num_tokens += block.tokens
+
+        return chunk_blocks, num_tokens, chunks
+
+    def _finalize_chunk(self, chunks: List[MarkdownChunk], blocks: List[BaseBlock], header: str, page_elements: dict):
+        """Create chunk with elements from the pages it spans"""
+        chunk_pages = set(block.page for block in blocks)
+        elements = []
+        for page_num in sorted(chunk_pages):
+            if page_num in page_elements:
+                elements.extend(page_elements[page_num])
+        chunks.append(MarkdownChunk(blocks=blocks, header=header, elements=elements))
+
+    def _create_chunk(self, chunks: List[MarkdownChunk], blocks: List[BaseBlock], header: str = None, page_elements: dict = None) -> Tuple[
+        List[MarkdownChunk], List[BaseBlock], int]:
+        """Creates a chunk and returns overlap blocks"""
+        page_elements = page_elements or {}
+        self._finalize_chunk(chunks, blocks, header, page_elements)
+
+        if not self.chunk_overlap:
+            return chunks, [], 0
+
+        overlap_tokens = 0
+        overlap_blocks = []
+
+        for block in reversed(blocks):
+            if block.block_type == "Text":
+                sentences = []
+
+                for sentence in reversed(block.sentences):
+
+                    if overlap_tokens + sentence.tokens > self.chunk_overlap:
+                        text_block = TextBlock.from_sentences(sentences=sentences, page=block.page)
+                        overlap_blocks.insert(0, text_block)
+                        return chunks, overlap_blocks, overlap_tokens
+
+                    else:
+                        sentences.insert(0, sentence)
+                        overlap_tokens += sentence.tokens
+
+            else:
+                if overlap_tokens + block.tokens > self.chunk_overlap:
+                    return chunks, overlap_blocks, overlap_tokens
+
+                else:
+                    overlap_blocks.insert(0, block)
+                    overlap_tokens += block.tokens
+
+        return chunks, [], 0