ai-parrot 0.17.2__cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

parrot/loaders/splitters/md.py

@@ -0,0 +1,228 @@
+from typing import List, Dict, Any
+import re
+from .base import BaseTextSplitter
+
+
+class MarkdownTextSplitter(BaseTextSplitter):
+    """
+    Text splitter that respects Markdown structure.
+
+    Features:
+    - Splits at markdown headers (maintaining hierarchy)
+    - Preserves code blocks
+    - Maintains lists and formatting
+    - Respects table structures
+    """
+
+    def __init__(
+        self,
+        chunk_size: int = 4000,
+        chunk_overlap: int = 200,
+        strip_headers: bool = False,
+        return_each_line: bool = False,
+        **kwargs
+    ):
+        """
+        Initialize MarkdownTextSplitter.
+
+        Args:
+            chunk_size: Maximum size per chunk (in characters)
+            chunk_overlap: Overlap between chunks
+            strip_headers: Whether to strip headers from chunks
+            return_each_line: Whether to return each line as a separate chunk
+        """
+        super().__init__(chunk_size, chunk_overlap, **kwargs)
+        self.strip_headers = strip_headers
+        self.return_each_line = return_each_line
+
+        # Markdown parsing patterns
+        self.header_pattern = re.compile(r'^(#{1,6})\s+(.+)$', re.MULTILINE)
+        self.code_block_pattern = re.compile(r'```.*?```', re.DOTALL)
+        self.list_pattern = re.compile(r'^[\s]*[-*+]\s+', re.MULTILINE)
+        self.table_pattern = re.compile(r'^[\s]*\|.*\|[\s]*$', re.MULTILINE)
+
+    def _count_tokens(self, text: str) -> int:
+        """Count tokens (approximation using words for markdown)"""
+        # Simple word-based counting for markdown
+        words = len(text.split())
+        return int(words * 1.3)  # Rough token approximation
+
+    def split_text(self, text: str) -> List[str]:
+        """Split markdown text respecting structure"""
+        if not text:
+            return []
+
+        if self.return_each_line:
+            return [line for line in text.split('\n') if line.strip()]
+
+        # First, identify markdown sections
+        sections = self._parse_markdown_sections(text)
+
+        # Then merge sections respecting chunk size
+        chunks = self._merge_markdown_sections(sections)
+
+        return chunks
+
+    def _parse_markdown_sections(self, text: str) -> List[Dict[str, Any]]:
+        """Parse markdown into hierarchical sections"""
+        lines = text.split('\n')
+        sections = []
+        current_section = {
+            'type': 'content',
+            'level': 0,
+            'header': '',
+            'content': [],
+            'start_line': 0
+        }
+
+        in_code_block = False
+        code_block_lang = ''
+
+        for i, line in enumerate(lines):
+            # Handle code blocks
+            if line.strip().startswith('```'):
+                if not in_code_block:
+                    # Starting code block
+                    in_code_block = True
+                    code_block_lang = line.strip()[3:]
+                    if current_section['content']:
+                        sections.append(current_section)
+                    current_section = {
+                        'type': 'code_block',
+                        'level': 0,
+                        'header': f'Code ({code_block_lang})',
+                        'content': [line],
+                        'start_line': i
+                    }
+                else:
+                    # Ending code block
+                    current_section['content'].append(line)
+                    sections.append(current_section)
+                    in_code_block = False
+                    current_section = {
+                        'type': 'content',
+                        'level': 0,
+                        'header': '',
+                        'content': [],
+                        'start_line': i + 1
+                    }
+                continue
+
+            if in_code_block:
+                current_section['content'].append(line)
+                continue
+
+            # Handle headers
+            header_match = self.header_pattern.match(line)
+            if header_match:
+                # Save current section if it has content
+                if current_section['content']:
+                    sections.append(current_section)
+
+                # Start new section
+                level = len(header_match.group(1))
+                header_text = header_match.group(2).strip()
+
+                current_section = {
+                    'type': 'section',
+                    'level': level,
+                    'header': header_text,
+                    'content': [] if self.strip_headers else [line],
+                    'start_line': i
+                }
+            else:
+                current_section['content'].append(line)
+
+        # Add final section
+        if current_section['content']:
+            sections.append(current_section)
+
+        return sections
+
+    def _merge_markdown_sections(self, sections: List[Dict[str, Any]]) -> List[str]:
+        """Merge sections respecting chunk size limits"""
+        if not sections:
+            return []
+
+        chunks = []
+        current_chunk_parts = []
+        current_size = 0
+
+        for section in sections:
+            section_text = '\n'.join(section['content'])
+            section_size = len(section_text)
+
+            # If section alone exceeds chunk size, split it
+            if section_size > self.chunk_size:
+                # Save current chunk if exists
+                if current_chunk_parts:
+                    chunks.append('\n\n'.join(current_chunk_parts))
+                    current_chunk_parts = []
+                    current_size = 0
+
+                # Split large section
+                split_sections = self._split_large_section(section_text)
+                chunks.extend(split_sections)
+                continue
+
+            # Check if adding this section exceeds chunk size
+            if current_size + section_size > self.chunk_size and current_chunk_parts:
+                # Save current chunk
+                chunks.append('\n\n'.join(current_chunk_parts))
+
+                # Start new chunk with overlap
+                overlap_parts = self._get_overlap_content(current_chunk_parts)
+                current_chunk_parts = overlap_parts + [section_text]
+                current_size = sum(len(part) for part in current_chunk_parts)
+            else:
+                current_chunk_parts.append(section_text)
+                current_size += section_size
+
+        # Add final chunk
+        if current_chunk_parts:
+            chunks.append('\n\n'.join(current_chunk_parts))
+
+        return [chunk for chunk in chunks if chunk.strip()]
+
+    def _split_large_section(self, text: str) -> List[str]:
+        """Split a large section that exceeds chunk size"""
+        # For very large sections, fall back to paragraph-based splitting
+        paragraphs = [p.strip() for p in text.split('\n\n') if p.strip()]
+
+        chunks = []
+        current_chunk = []
+        current_size = 0
+
+        for paragraph in paragraphs:
+            para_size = len(paragraph)
+
+            if current_size + para_size > self.chunk_size and current_chunk:
+                chunks.append('\n\n'.join(current_chunk))
+                current_chunk = [paragraph]
+                current_size = para_size
+            else:
+                current_chunk.append(paragraph)
+                current_size += para_size
+
+        if current_chunk:
+            chunks.append('\n\n'.join(current_chunk))
+
+        return chunks
+
+    def _get_overlap_content(self, parts: List[str]) -> List[str]:
+        """Get content for overlap between chunks"""
+        if not parts or self.chunk_overlap == 0:
+            return []
+
+        # Take last parts that fit in overlap size
+        overlap_parts = []
+        overlap_size = 0
+
+        for part in reversed(parts):
+            if overlap_size + len(part) <= self.chunk_overlap:
+                overlap_parts.insert(0, part)
+                overlap_size += len(part)
+            else:
+                break
+
+        return overlap_parts

parrot/loaders/splitters/token.py

@@ -0,0 +1,143 @@
+from typing import List, Callable, Optional, Any, Union
+from .base import BaseTextSplitter
+
+
+class TokenTextSplitter(BaseTextSplitter):
+    """
+    Text splitter that splits based on token count using various tokenizers.
+
+    Supports:
+    - OpenAI tiktoken tokenizers (gpt-3.5-turbo, gpt-4, etc.)
+    - Hugging Face transformers tokenizers
+    - Custom tokenization functions
+    """
+
+    def __init__(
+        self,
+        chunk_size: int = 4000,
+        chunk_overlap: int = 200,
+        model_name: str = "gpt-3.5-turbo",
+        encoding_name: Optional[str] = None,
+        tokenizer: Optional[Any] = None,
+        tokenizer_function: Optional[Callable[[str], List[str]]] = None,
+        allowed_special: Union[str, set] = "all",
+        disallowed_special: Union[str, set, List] = (),
+        **kwargs
+    ):
+        """
+        Initialize TokenTextSplitter.
+
+        Args:
+            chunk_size: Maximum number of tokens per chunk
+            chunk_overlap: Number of tokens to overlap between chunks
+            model_name: Name of the model for tiktoken (e.g., 'gpt-3.5-turbo', 'gpt-4')
+            encoding_name: Specific encoding name for tiktoken (e.g., 'cl100k_base')
+            tokenizer: Custom tokenizer object (Hugging Face or other)
+            tokenizer_function: Custom function that takes text and returns tokens
+            allowed_special: Special tokens allowed during encoding
+            disallowed_special: Special tokens that raise errors during encoding
+        """
+        super().__init__(chunk_size, chunk_overlap, **kwargs)
+
+        self.model_name = model_name
+        self.encoding_name = encoding_name
+        self.allowed_special = allowed_special
+        self.disallowed_special = disallowed_special
+
+        # Initialize tokenizer
+        if tokenizer_function:
+            self._tokenize = tokenizer_function
+            self._encode = lambda text: tokenizer_function(text)
+            self._decode = lambda tokens: " ".join(tokens)
+        elif tokenizer:
+            # Hugging Face tokenizer
+            self._tokenizer = tokenizer
+            self._tokenize = self._hf_tokenize
+            self._encode = self._hf_encode
+            self._decode = self._hf_decode
+        else:
+            # Use tiktoken
+            import tiktoken
+            try:
+                if encoding_name:
+                    self._enc = tiktoken.get_encoding(encoding_name)
+                else:
+                    self._enc = tiktoken.encoding_for_model(model_name)
+
+                self._tokenize = self._tiktoken_tokenize
+                self._encode = self._tiktoken_encode
+                self._decode = self._tiktoken_decode
+            except Exception as e:
+                raise ValueError(f"Could not load tokenizer for model {model_name}: {e}")
+
+    def _tiktoken_tokenize(self, text: str) -> List[str]:
+        """Tokenize using tiktoken and return string tokens"""
+        tokens = self._enc.encode(
+            text,
+            allowed_special=self.allowed_special,
+            disallowed_special=self.disallowed_special
+        )
+        return [self._enc.decode([token]) for token in tokens]
+
+    def _tiktoken_encode(self, text: str) -> List[int]:
+        """Encode text to token IDs using tiktoken"""
+        return self._enc.encode(
+            text,
+            allowed_special=self.allowed_special,
+            disallowed_special=self.disallowed_special
+        )
+
+    def _tiktoken_decode(self, tokens: List[int]) -> str:
+        """Decode token IDs to text using tiktoken"""
+        return self._enc.decode(tokens)
+
+    def _hf_tokenize(self, text: str) -> List[str]:
+        """Tokenize using Hugging Face tokenizer"""
+        return self._tokenizer.tokenize(text)
+
+    def _hf_encode(self, text: str) -> List[int]:
+        """Encode text using Hugging Face tokenizer"""
+        return self._tokenizer.encode(text, add_special_tokens=False)
+
+    def _hf_decode(self, tokens: List[int]) -> str:
+        """Decode tokens using Hugging Face tokenizer"""
+        return self._tokenizer.decode(tokens, skip_special_tokens=True)
+
+    def _count_tokens(self, text: str) -> int:
+        """Count tokens in text"""
+        return len(self._encode(text))
+
+    def split_text(self, text: str) -> List[str]:
+        """Split text based on token count"""
+        if not text:
+            return []
+
+        # Encode the entire text
+        tokens = self._encode(text)
+
+        if len(tokens) <= self.chunk_size:
+            return [text]
+
+        chunks = []
+        start_idx = 0
+
+        while start_idx < len(tokens):
+            # Determine end index for this chunk
+            end_idx = min(start_idx + self.chunk_size, len(tokens))
+
+            # Get tokens for this chunk
+            chunk_tokens = tokens[start_idx:end_idx]
+
+            # Decode back to text
+            chunk_text = self._decode(chunk_tokens)
+            chunks.append(chunk_text)
+
+            # Move start index with overlap
+            if end_idx == len(tokens):
+                break
+
+            start_idx = end_idx - self.chunk_overlap
+            if start_idx < 0:
+                start_idx = 0
+
+        return chunks

parrot/loaders/txt.py

@@ -0,0 +1,26 @@
+from typing import List
+from pathlib import Path
+from .abstract import AbstractLoader
+from .files.text import TextFile
+
+
+class TextLoader(AbstractLoader):
+    """
+    Loader for Text-based Files.
+    """
+    extensions: List[str] = ['.txt', '.text', '.md', '.markdown', '.rd']
+
+    async def _load(self, path: Path, **kwargs) -> list:
+        """
+        Load a TXT file.
+
+        Args:
+            path (Path): The path to the TXT file.
+
+        Returns:
+            list: A list of Langchain Documents.
+        """
+        async with TextFile(path) as file:
+            content = await file.read()
+            return self.create_document(content, path)
+        return []

parrot/loaders/video.py

@@ -0,0 +1,89 @@
+from typing import List
+from abc import abstractmethod
+from pathlib import Path
+import subprocess
+from ..stores.models import Document
+from .basevideo import BaseVideoLoader
+
+
+class VideoLoader(BaseVideoLoader):
+    """
+    Generating Video transcripts from URL Videos.
+    """
+    def download_video(self, url: str, path: str) -> Path:
+        """
+        Downloads a video from a URL using yt-dlp.
+
+        Args:
+            video_url (str): The URL of the video to download.
+            output_path (str): The directory where the video will be saved.
+        """
+        try:
+            command = [
+                "yt-dlp",
+                "--get-filename",
+                "-o",
+                str(path / "%(title)s.%(ext)s"),
+                url
+            ]
+            result = subprocess.run(command, check=True, stdout=subprocess.PIPE, text=True)
+        except Exception as e:
+            try:
+                command = [
+                    "yt-dlp",
+                    "--get-filename",
+                    url
+                ]
+                result = subprocess.run(
+                    command,
+                    check=True,
+                    stdout=subprocess.PIPE,
+                    text=True
+                )
+            except Exception as e:
+                raise ValueError(
+                    f"Unable to Download Video: {e}"
+                )
+        try:
+            filename = result.stdout.strip()  # Remove any trailing newline characters
+            print('FILENAME > ', filename)
+            file_path = path.joinpath(filename)
+            if file_path.exists():
+                print(f"Video already downloaded: {filename}")
+                return file_path
+            print(f"Downloading video: {filename}")
+            # after extracted filename, download the video
+            command = [
+                "yt-dlp",
+                url,
+                "-o",
+                str(path / "%(title)s.%(ext)s")
+            ]
+            subprocess.run(command, check=True)
+            return file_path
+        except subprocess.CalledProcessError as e:
+            print(f"Error downloading video: {e}")
+
+
+    async def _load(self, source: str, **kwargs) -> List[Document]:
+        documents = []
+        transcript = None
+        video_title = source
+        if isinstance(source, dict):
+            path = list(source.keys())[0]
+            parts = source[path]
+            if isinstance(parts, str):
+                video_title = parts
+            elif isinstance(parts, dict):
+                video_title = parts['title']
+        docs = await self.load_video(source, video_title, transcript)
+        documents.extend(docs)
+        # return documents
+        return documents
+
+    @abstractmethod
+    async def load_video(self, url: str, video_title: str, transcript: str) -> list:
+        pass
+
+    def parse(self, source):
+        pass