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

parrot/loaders/pdfchapters.py

@@ -0,0 +1,142 @@
+from collections.abc import Callable
+from typing import Any, Optional, List, Union
+from pathlib import PurePath, Path
+from io import StringIO
+import fitz  # PyMuPDF
+from langchain.docstore.document import Document
+from langchain.text_splitter import (
+    RecursiveCharacterTextSplitter
+)
+from .basepdf import BasePDF
+
+
+class PDFChapterLoader(BasePDF):
+    """
+    Preserving Chapter Structure from PDF files.
+    """
+    def __init__(
+        self,
+        path: PurePath,
+        tokenizer: Callable[..., Any] = None,
+        text_splitter: Callable[..., Any] = None,
+        source_type: str = 'pdf',
+        language: str = "eng",
+        **kwargs
+    ):
+        super().__init__(
+            path=path,
+            tokenizer=tokenizer,
+            text_splitter=text_splitter,
+            source_type=source_type,
+            language=language,
+            **kwargs
+        )
+        # Which Font is used for titles (Chapter separation)
+        self.title_font: list = kwargs.get('title_font', 'Calibri-Bold')
+        if not text_splitter:
+            self.text_splitter = RecursiveCharacterTextSplitter.from_huggingface_tokenizer(
+                self.tokenizer,
+                chunk_size=2000,
+                chunk_overlap=100,
+                add_start_index=True,  # If `True`, includes chunk's start index in metadata
+                strip_whitespace=True,  # strips whitespace from the start and end
+                separators=["\n\n", "\n", "\r\n", "\r", "\f", "\v", "\x0b", "\x0c"],
+            )
+
+    def eval_title(self, title_font: str) -> bool:
+        """
+        Check if the font is a title font.
+
+        Args:
+            title_font (str): The font to check.
+
+        Returns:
+            bool: True if the font is a title font.
+        """
+        return 'Bold' in title_font or title_font == self.title_font
+
+    def _load_pdf(self, path: PurePath, **kwargs):
+        """
+        Open a PDF file.
+
+        Args:
+            path (PurePath): The path to the PDF file.
+
+        Returns:
+            pdfplumber.PDF: The PDF object.
+        """
+        pdf = fitz.open(path)
+        self.logger.info(f"Loading PDF file: {path}")
+        chapters = []
+        current_chapter_text = ''
+        current_chapter_title = ''
+        current_chapter_page = None
+        chapter_titles = set()  # Keep track of unique chapter titles
+        for page_num in range(len(pdf)):
+            page = pdf.load_page(page_num)
+            blocks = page.get_text("dict")["blocks"]
+            page_number = page_num + 1
+            metadata = {
+                "url": '',
+                "index": f"{path.name} #{page_number}",
+                "source": f"{path.name} #{page_number}",
+                "filename": path.name,
+                "source_type": self._source_type,
+                "type": "pdf",
+                "question": "",
+                "answer": "",
+                "summary": '',
+                "document_meta": {
+                    "page_number": page_num,
+                    # **pdf.metadata
+                }
+            }
+            for b in blocks:
+                if b['type'] == 0:  # Text block
+                    block_text = ''
+                    for line in b["lines"]:
+                        for span in line["spans"]:
+                            block_text += span['text']  # Accumulate text within the block
+
+                    # Check if the block text is a title by examining the font
+                    if any(self.eval_title(span['font']) for line in b["lines"] for span in line["spans"]):
+                        title = block_text.strip()
+                        if title not in chapter_titles:
+                            # Save the current chapter if it's not empty and start a new one
+                            if current_chapter_text.strip() and current_chapter_text.strip() != current_chapter_title.strip():
+                                chapters.append({
+                                    'chapter': current_chapter_title,
+                                    'content': current_chapter_text.strip(),
+                                    'page': current_chapter_page,
+                                    'meta': metadata
+                                })
+                            current_chapter_title = f"**{title}**: "
+                            current_chapter_page = page_num + 1
+                            current_chapter_text = current_chapter_title
+                            chapter_titles.add(title)
+                        else:
+                            # Continue appending to the existing chapter
+                            current_chapter_text += block_text
+                    else:
+                        # Continue appending text to the current chapter
+                        current_chapter_text += block_text
+
+            # Add a newline after processing each block, if not a chapter title
+            if not block_text.strip().startswith(current_chapter_title):
+                current_chapter_text += "\n"
+
+        # Save the last chapter if it exists and it's not just the title
+        if current_chapter_text.strip() and current_chapter_text.strip() != current_chapter_title.strip():
+            chapters.append({
+                'chapter': current_chapter_title,
+                'content': current_chapter_text.strip(),
+                'page': current_chapter_page,
+                'meta': metadata
+            })
+        documents = []
+        for chapter in chapters:
+            documents.append(Document(
+                page_content=chapter['content'],
+                metadata=chapter['meta']
+            ))
+        return documents

parrot/loaders/pdffn.py

@@ -0,0 +1,112 @@
+from collections.abc import Callable
+from typing import Any, Optional, List, Union
+from pathlib import PurePath, Path
+from io import StringIO
+import fitz  # PyMuPDF
+from langchain.docstore.document import Document
+from langchain.text_splitter import (
+    RecursiveCharacterTextSplitter
+)
+from .basepdf import BasePDF
+
+
+class PDFFnLoader(BasePDF):
+    """
+    Loading a PDF with including function processing.
+    """
+    def __init__(
+        self,
+        path: PurePath,
+        tokenizer: Callable[..., Any] = None,
+        text_splitter: Union[None, Callable[..., Any]] = None,
+        source_type: str = 'pdf',
+        language: str = "eng",
+        **kwargs
+    ):
+        table_settings = kwargs.pop('table_settings', {})
+        super().__init__(
+            path=path,
+            tokenizer=tokenizer,
+            text_splitter=text_splitter,
+            source_type=source_type,
+            language=language,
+            **kwargs
+        )
+        if not text_splitter:
+            self.text_splitter = RecursiveCharacterTextSplitter.from_huggingface_tokenizer(
+                self.tokenizer,
+                chunk_size=2000,
+                chunk_overlap=100,
+                add_start_index=True,  # If `True`, includes chunk's start index in metadata
+                strip_whitespace=True,  # strips whitespace from the start and end
+                separators=["\n\n", "\n", "\r\n", "\r", "\f", "\v", "\x0b", "\x0c"],
+            )
+        self.table_settings = {
+            "vertical_strategy": "lines",
+            "horizontal_strategy": "text",
+            "intersection_x_tolerance": 5,
+            "intersection_y_tolerance": 5,
+            "edge_min_length": 10,
+        }
+        # Define settings for Fitz Table Processing
+        self.table_settings = {**self.table_settings, **table_settings}
+
+    def set_metadata(self, path, page, page_number, **kwargs) -> dict:
+        n = page_number + 1
+        return {
+            "url": '',
+            "index": f"{path.name} #{page_number}",
+            "source": f"{path.name} #{page_number}",
+            "filename": path.name,
+            "source_type": self._source_type,
+            "type": "pdf",
+            "question": "",
+            "answer": "",
+            "summary": '',
+            "document_meta": {
+                "page_number": n,
+                **kwargs
+            }
+        }
+
+    def processing_table(self, table, table_idx, page, **kwargs) -> dict:
+        df = table.to_pandas()  # convert to pandas DataFrame
+        df = df.dropna(axis=1, how='all')
+        df = df.dropna(how='all', axis=0)  # Drop empty rows
+        table_data = []
+        # Extract text from each cell
+        for row_idx in range(table.row_count):
+            for col_idx in range(table.column_count):
+                cell = table[row_idx][col_idx]
+                print('CELL ', cell)
+                print('---------')
+                cell_text = cell.get_text("text", flags=fitz.TEXTFLAGS_HTML)
+                print(cell_text)
+
+        return table_data
+
+    def _load_pdf(self, path: PurePath, **kwargs):
+        """
+        Open a PDF file.
+
+        Args:
+            path (PurePath): The path to the PDF file.
+
+        Returns:
+            fitz.PDF: The PDF object.
+        """
+        pdf = fitz.open(path)
+        self.logger.info(f"Loading PDF file: {path}")
+        for page_num in range(len(pdf)):
+            # Will extract first the table and second the block of texts
+            page = pdf.load_page(page_num)
+            parts = page.get_text("dict", flags=fitz.TEXTFLAGS_HTML)
+            # print('PARTS ', parts)
+            blocks = page.get_text("dict")["blocks"]
+            # print('BLOCKS >', blocks)
+            metadata = self.set_metadata(path, page, page_num)
+            # print('META > ', metadata)
+            tables = page.find_tables(**self.table_settings)
+            for tab_idx, table in enumerate(tables):
+                table_data = self.processing_table(table, tab_idx, page)
+        return []

parrot/loaders/pdfimages.py

@@ -0,0 +1,207 @@
+from collections.abc import Callable
+from pathlib import Path, PurePath
+from typing import Any
+import fitz
+from pdf4llm import to_markdown
+from PIL import Image
+from langchain.docstore.document import Document
+from langchain.text_splitter import MarkdownTextSplitter
+from transformers import (
+    AutoTokenizer,
+    AutoProcessor,
+    LlavaForConditionalGeneration,
+    pipeline,
+    BitsAndBytesConfig
+)
+import torch
+from .basepdf import BasePDF
+
+
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float16
+)
+
+
+class PDFImageLoader(BasePDF):
+    """
+    Loader for PDF files.
+    """
+    default_prompt: str = "<|user|>\n<image>\nExplain this schematic diagram or technical installation instructions and wire diagrams, please be detailed about descriptions of steps:<|end|>\n<|assistant|>\n"
+    def __init__(
+        self,
+        path: PurePath,
+        tokenizer: Callable[..., Any] = None,
+        text_splitter: Callable[..., Any] = None,
+        source_type: str = 'pdf',
+        language: str = "eng",
+        **kwargs
+    ):
+        super().__init__(
+            path=path,
+            tokenizer=tokenizer,
+            text_splitter=text_splitter,
+            source_type=source_type,
+            language=language,
+            **kwargs
+        )
+        self._image_model = kwargs.get('image_model', 'llava-hf/llava-v1.6-vicuna-7b-hf')
+        self._task = kwargs.get('task', 'image-to-text')
+        self._max_tokens = kwargs.get('max_tokens', 600)
+        # Loading the model with low CPU memory usage
+        # model = LlavaForConditionalGeneration.from_pretrained(
+        #     self._image_model,
+        #     quantization_config=quantization_config,
+        #     device_map="auto",
+        #     torch_dtype=torch.float16,
+        #     low_cpu_mem_usage=True
+        # )
+        # # Load the processor
+        # processor = AutoProcessor.from_pretrained(self._image_model, use_fast=True)
+        self._pipeline = pipeline(
+            self._task,
+            model=self._image_model,
+            # tokenizer=processor.tokenizer,
+            # image_processor=processor.image_processor,
+            model_kwargs={"quantization_config": quantization_config},
+            # device=self._device,
+            max_new_tokens=self._max_tokens,
+            # low_cpu_mem_usage=True,
+            use_fast=True
+        )
+        # default prompt
+        self._prompt = kwargs.get('prompt', self.default_prompt)
+        # Markdown Splitter
+        self._splitter = MarkdownTextSplitter(
+            chunk_size = self._chunk_size,
+            chunk_overlap=10
+        )
+
+    def pixmap_to_pil_image(self, pix):
+        """Converts a PyMuPDF Pixmap object to a PIL Image"""
+        return Image.frombytes(
+            "RGB",
+            [pix.width, pix.height],
+            pix.samples
+        )
+
+    def _load_pdf(self, path: Path) -> list:
+        """
+        Load a PDF file as Images.
+
+        Args:
+            path (Path): The path to the PDF file.
+
+        Returns:
+            list: A list of Langchain Documents.
+        """
+        if self._check_path(path):
+            self.logger.info(f"Loading PDF file: {path}")
+            pdf = fitz.open(str(path))  # Open the PDF file
+            docs = []
+            try:
+                # get markdown for all pages and saved separately
+                md_text = to_markdown(pdf)
+                try:
+                    summary = self.get_summary_from_text(md_text)
+                except Exception:
+                    summary = ''
+                metadata = {
+                    "url": '',
+                    "idx": str(path.name),
+                    "filename": str(path.name),
+                    "source": str(path.name),
+                    "type": 'pdf',
+                    "question": '',
+                    "answer": '',
+                    "data": {},
+                    "summary": summary,
+                    "source_type": self._source_type,
+                    "document_meta": {
+                        "title": pdf.metadata.get("title", ""),
+                        "author": pdf.metadata.get("author", ""),
+                    }
+                }
+                for idx, chunk in enumerate(self._splitter.split_text(md_text)):
+                    _info = {
+                        "index": f"{idx}",
+                        **metadata
+                    }
+                    docs.append(
+                        Document(
+                            page_content=chunk,
+                            metadata=_info
+                        )
+                    )
+            except (IndexError, ValueError) as exc:
+                self.logger.warning(
+                    f"There is no text data to load on {path.name}: {exc}"
+                )
+            # Then, processing the pages one by one as Images:
+            file_name = path.stem.replace(" ", "_").replace(".", "_")
+            for page_number in range(pdf.page_count):
+                page_num = page_number + 1
+                self.logger.notice(
+                    f"Processing PDF {path} on Page {page_num}"
+                )
+                page = pdf[page_number]
+                pix = page.get_pixmap(colorspace=fitz.csRGB, alpha=False)
+                zoom_x = 2.0  # horizontal zoom
+                zoom_y = 2.0  # vertical zoom
+                mat = fitz.Matrix(zoom_x, zoom_y)  # zoom factor 2 in each dimension
+                img_stream = self.pixmap_to_pil_image(pix)
+                url = ''
+                img_name = f'image_{file_name}_{page_num}.png'
+                if self.save_images is True:
+                    img_path = self.save_image(
+                        img_stream,
+                        img_name,
+                        self._imgdir
+                    )
+                    url = f'/static/images/{img_name}'
+                # extracting features and explanations:
+                outputs = self._pipeline(
+                    img_stream,
+                    prompt=self._prompt,
+                    generate_kwargs={"max_new_tokens": self._max_tokens}
+                )
+                documents = []
+                for idx, output in enumerate(outputs):
+                    generated_text = output['generated_text']
+                    # Split using the special tokens, if available
+                    split_text = generated_text.split("<|assistant|>")
+                    prompt_text = split_text[0].replace("<|prompt|>", "").strip() if "<|prompt|>" in generated_text else ""
+                    response_text = split_text[1].strip() if len(split_text) > 1 else ""
+                    # Attach the image using Markdown syntax
+                    image_markdown = f"\n\n![Image]({url})\n"
+                    response_text += image_markdown
+                    _meta = {
+                        "url": f"{url}",
+                        "filename": str(path.name),
+                        "index": f"Page {page_num}, part: {idx}",
+                        "source": str(path.name),
+                        "type": 'pdf',
+                        "question": prompt_text,
+                        "answer": '',
+                        "data": {},
+                        "summary": '',
+                        "source_type": self._source_type,
+                        "document_meta": {
+                            "page": f"Page {page}",
+                            "image": f"{img_name}",
+                            "url": f"{url}"
+                        }
+                    }
+                    documents.append(
+                        Document(
+                            page_content=response_text,
+                            metadata=_meta
+                        )
+                    )
+            return docs + documents
+
+    def load(self) -> list:
+        try:
+            return super().load()
+        finally:
+            self._pipeline = None

parrot/loaders/pdfmark.py

@@ -0,0 +1,88 @@
+from typing import Any
+from collections.abc import Callable
+from pathlib import Path, PurePath
+import fitz
+from pdf4llm import to_markdown
+from langchain.docstore.document import Document
+from langchain.text_splitter import MarkdownTextSplitter
+from .basepdf import BasePDF
+
+class PDFMarkdownLoader(BasePDF):
+    """
+    Loader for PDF files converted content to markdown.
+    """
+
+    def __init__(
+        self,
+        path: PurePath,
+        tokenizer: Callable[..., Any] = None,
+        text_splitter: Callable[..., Any] = None,
+        source_type: str = 'pdf',
+        language: str = "eng",
+        **kwargs
+    ):
+        super().__init__(
+            path=path,
+            tokenizer=tokenizer,
+            text_splitter=text_splitter,
+            source_type=source_type,
+            language=language,
+            **kwargs
+        )
+        self._splitter = MarkdownTextSplitter(chunk_size = 1024, chunk_overlap=10)
+
+    def _load_pdf(self, path: Path) -> list:
+        """
+        Load a PDF file using the PDFMiner library.
+
+        Args:
+            path (Path): The path to the PDF file.
+
+        Returns:
+            list: A list of Langchain Documents.
+        """
+        if self._check_path(path):
+            self.logger.info(f"Loading PDF file: {path}")
+            docs = []
+            pdf = fitz.open(str(path))
+            md_text = to_markdown(pdf) # get markdown for all pages
+            try:
+                summary = self.get_summary_from_text(md_text)
+            except Exception:
+                summary = ''
+            metadata = {
+                "url": '',
+                "filename": path.name,
+                # "index": f"{path.name}",
+                "source": str(path.name),
+                "type": 'pdf',
+                "question": '',
+                "answer": '',
+                "data": {},
+                "summary": summary,
+                "source_type": self._source_type,
+                "document_meta": {
+                    "title": pdf.metadata.get("title", ""),
+                    # "subject": pdf.metadata.get("subject", ""),
+                    # "keywords": pdf.metadata.get("keywords", ""),
+                    "creationDate": pdf.metadata.get("creationDate", ""),
+                    # "modDate": pdf.metadata.get("modDate", ""),
+                    # "producer": pdf.metadata.get("producer", ""),
+                    # "creator": pdf.metadata.get("creator", ""),
+                    "author": pdf.metadata.get("author", ""),
+                }
+            }
+            for idx, chunk in enumerate(self._splitter.split_text(md_text)):
+                _info = {
+                    "index": f"{idx}",
+                    **metadata
+                }
+                docs.append(
+                    Document(
+                        page_content=chunk,
+                        metadata=_info
+                    )
+                )
+            return docs
+        else:
+            return []

parrot/loaders/pdftables.py

@@ -0,0 +1,145 @@
+from collections.abc import Callable
+from typing import Any, Optional, List
+from pathlib import Path, PurePath
+from io import StringIO
+import fitz
+import pandas as pd
+from langchain.docstore.document import Document
+from .basepdf import BasePDF
+
+
+class PDFTablesLoader(BasePDF):
+    """
+    Loader for Tables in PDF Files.
+    """
+    _extension = ['.pdf']
+
+    def __init__(
+        self,
+        path: PurePath,
+        tokenizer: Callable[..., Any] = None,
+        text_splitter: Callable[..., Any] = None,
+        source_type: str = 'pdf',
+        language: str = "eng",
+        table_settings: dict = {},
+        **kwargs
+    ):
+        super().__init__(
+            path,
+            tokenizer,
+            text_splitter,
+            source_type,
+            language=language,
+            **kwargs
+        )
+        # Table Settings:
+        self.table_settings = {
+            #"vertical_strategy": "text",
+            # "horizontal_strategy": "text",
+            "intersection_x_tolerance": 5,
+            "intersection_y_tolerance": 5
+        }
+        if table_settings:
+            self.table_settings.update(table_settings)
+        self._skiprows = kwargs.pop('skiprows', None)
+
+    def unique_columns(self, df: pd.DataFrame) -> pd.DataFrame:
+        """
+        Rename duplicate columns in the DataFrame to ensure they are unique.
+
+        Args:
+            df (pd.DataFrame): The DataFrame with potential duplicate column names.
+
+        Returns:
+            pd.DataFrame: A DataFrame with unique column names.
+        """
+        seen = {}
+        new_columns = []
+        for col in df.columns:
+            new_col = col
+            count = seen.get(col, 0)
+            while new_col in new_columns:
+                count += 1
+                new_col = f"{col}_{count}"
+            new_columns.append(new_col)
+            seen[col] = count
+        df.columns = new_columns
+        return df
+
+    def get_markdown(self, df: pd.DataFrame) -> str:
+        """
+        Convert a DataFrame to a Markdown string.
+
+        Args:
+            df (pd.DataFrame): The DataFrame to convert.
+
+        Returns:
+            str: The JSON string.
+        """
+        buffer = StringIO()
+        df = self.unique_columns(df)
+        df.to_markdown(buffer)
+        buffer.seek(0)
+        return buffer.getvalue()
+
+    def parse_table(self, table_idx, table, page_number, path) -> pd.DataFrame:
+        df = table.to_pandas()  # convert to pandas DataFrame
+        df = df.dropna(axis=1, how='all')
+        df = df.dropna(how='all', axis=0)  # Drop empty rows
+        page = page_number + 1
+        table_meta = {
+            "url": '',
+            "source": f"{path.name} Page.#{page} Table.#{table_idx}",
+            "filename": path.name,
+            "index": f"{path.name}:Table:{table_idx}",
+            "question": '',
+            "answer": '',
+            "type": 'table',
+            "data": {},
+            "summary": '',
+            "document_meta": {
+                "table_index": table_idx,
+                "table_shape": df.shape,
+                "table_columns": df.columns.tolist(),
+                "description": f"Extracted from Page.#{page}."
+            },
+            "source_type": self._source_type
+        }
+        return df, table_meta
+
+    def _load_pdf(self, path: Path) -> list:
+        """
+        Load a PDF file using the Fitz library.
+
+        Args:
+            path (Path): The path to the PDF file.
+
+        Returns:
+            list: A list of Langchain Documents.
+        """
+        if self._check_path(path):
+            self.logger.info(f"Loading PDF file: {path}")
+            pdf = fitz.open(str(path))  # Open the PDF file
+            docs = []
+            for page_number in range(pdf.page_count):
+                page = pdf[page_number]
+                try:
+                    tabs = page.find_tables(**self.table_settings)
+                    for tab_idx, tab in enumerate(tabs):
+                        df, _meta = self.parse_table(tab_idx, tab, page_number, path)
+                        ## Sample information:
+                        print('::: Printing Table Information === ')
+                        print(df)
+                        print("::: Printing Column Information === ")
+                        for column, t in df.dtypes.items():
+                            print(column, "->", t, "->", df[column].iloc[0])
+                        # convert into markdown:
+                        txt = df.to_markdown()
+                        if txt:
+                            docs.append(
+                                Document(page_content=txt, metadata=_meta)
+                            )
+                except Exception as exc:
+                    print(exc)
+                    continue
+            return docs