ai-parrot 0.1.0__cp311-cp311-manylinux_2_28_x86_64.whl

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

parrot/loaders/ppt.py

@@ -0,0 +1,30 @@
+from pathlib import PurePath
+from langchain_community.document_loaders import (
+    UnstructuredPowerPointLoader
+)
+from .abstract import AbstractLoader
+
+
+class PPTXLoader(AbstractLoader):
+    """
+    Loader for PPTX files.
+    """
+    _extension: list = ['.pptx']
+
+    def load(self, path: PurePath) -> list:
+        if self._check_path(path):
+            docs = []
+            self.logger.info(f"Loading PPTX file: {path}")
+            ppt_loader = UnstructuredPowerPointLoader(
+                file_path=str(path)
+            )
+            docs += ppt_loader.load()
+            for doc in docs:
+                doc.metadata['source_type'] = self._source_type
+            # Split the documents into chunks
+            return self.split_documents(docs)
+        else:
+            return []
+
+    def parse(self, source):
+        pass

parrot/loaders/qa.py

@@ -0,0 +1,81 @@
+
+from pathlib import Path, PurePath
+from typing import Any
+from collections.abc import Callable
+import pandas as pd
+from langchain.docstore.document import Document
+from .abstract import AbstractLoader
+
+
+class QAFileLoader(AbstractLoader):
+    """
+    Question and Answers File based on Excel.
+    """
+    _extension = ['.xlsx']
+    chunk_size = 768
+
+    def __init__(
+        self,
+        path: PurePath,
+        tokenizer: Callable[..., Any] = None,
+        text_splitter: Callable[..., Any] = None,
+        source_type: str = 'QA',
+        columns: list = ['Question', 'Answer'],
+        **kwargs
+    ):
+        super().__init__(tokenizer, text_splitter, source_type, **kwargs)
+        self.path = path
+        self._columns = columns
+        if isinstance(path, str):
+            self.path = Path(path).resolve()
+        if self.path.is_dir():
+            raise ValueError(
+                f"Currently only accepting single Files."
+            )
+
+    def _load_document(self, path: PurePath) -> list:
+        if path.exists():
+            print('Load QA Excel File: ', path)
+            df = pd.read_excel(path)
+            q = self._columns[0]
+            a = self._columns[1]
+            docs = []
+            for idx, row in df.iterrows():
+                # Question Document
+                doc = Document(
+                    page_content=f"**Question:** {row[q]}: **Answer:** {row[a]}",
+                    metadata={
+                        "url": '',
+                        "index": f"{path.name} #{idx}",
+                        "source": f"{path.name} Row.#{idx}",
+                        "filename": f"{path.name}",
+                        "question": row[q],
+                        "answer": row[a],
+                        "page_number": idx,
+                        "source_type": self._source_type,
+                        "type": "QA",
+                        "summary": f"Question: {row[q]}?: **{row[a]}**",
+                        "document_meta": {
+                            "question": row[q],
+                            "answer": row[a],
+                        }
+                    }
+                )
+                docs.append(doc)
+            return docs
+        return []
+
+    def load(self, **kwargs) -> list:
+        """
+        Load Chapters from a PDF file.
+
+        Returns:
+            list: A list of Langchain Documents.
+        """
+        if self.path.is_file():
+            documents = self._load_document(path=self.path, **kwargs)
+            # after all documents are retrieved, procesed and stored
+            return self.split_documents(documents)
+
+    def parse(self, source):
+        pass

parrot/loaders/repo.py

@@ -0,0 +1,103 @@
+from pathlib import PurePath
+from langchain_core.document_loaders.blob_loaders import Blob
+from langchain_community.document_loaders.generic import GenericLoader
+from langchain_community.document_loaders.parsers import LanguageParser
+from langchain_community.document_loaders import (
+    DirectoryLoader,
+    TextLoader,
+    JSONLoader
+)
+from langchain_text_splitters import Language
+from langchain.text_splitter import (
+    RecursiveCharacterTextSplitter
+)
+from .abstract import AbstractLoader
+
+
+class RepositoryLoader(AbstractLoader):
+    """Repository (Code Directory) loader.
+    """
+    exclude_paths: list = [
+        ".venv/**",
+        ".venv/**/**/*",
+        ".git/**",
+        "node_modules/**",
+        "build/**",
+        "dist/**",
+        "templates/**",
+        "tmp/**"
+    ]
+
+    def load(self, path: PurePath, lang: str = 'python', excludes: list = []) -> list:
+        """
+        Load data from a repository and return it as a Langchain Document.
+        """
+        if isinstance(path, str):
+            path = PurePath(path)
+        if excludes:
+            self.exclude_paths += excludes
+        excludes_path = [
+            str(path.joinpath(p).resolve()) for p in self.exclude_paths
+        ]
+        if lang == 'python':
+            parser = LanguageParser(language=Language.PYTHON, parser_threshold=100)
+            splitter = RecursiveCharacterTextSplitter.from_language(
+                language=Language.PYTHON, chunk_size=1024, chunk_overlap=200
+            )
+            suffixes = [".py", ".pyx"]
+            glob = "**/[!.]*.py?"
+        elif lang == 'javascript':
+            parser = LanguageParser(language=Language.JS, parser_threshold=100)
+            splitter = RecursiveCharacterTextSplitter.from_language(
+                language=Language.JS, chunk_size=1024, chunk_overlap=200
+            )
+            suffixes = [".js", ".jsx", ".json", ".ts", ".tsx"]
+        elif lang == 'typescript':
+            parser = LanguageParser(language=Language.TS, parser_threshold=100)
+            splitter = RecursiveCharacterTextSplitter.from_language(
+                language=Language.TS, chunk_size=1024, chunk_overlap=200
+            )
+            suffixes = [".js", ".jsx", ".json", ".ts", ".tsx"]
+        elif lang == 'json':
+            loader = DirectoryLoader(
+                path,
+                glob="**/*.json",
+                show_progress=True,
+                exclude=excludes_path,
+                silent_errors=True,
+                recursive=True,
+                # loader_cls=TextLoader,
+                loader_cls=JSONLoader,
+                loader_kwargs={
+                    'jq_schema': '.',
+                    'text_content': False
+                }
+            )
+            docs = loader.load()
+            for doc in docs:
+                doc.metadata['url'] = ''
+                doc.metadata['source_type'] = self._source_type
+                doc.metadata['language'] = lang
+            return self.text_splitter.split_documents(docs)
+        else:
+            raise ValueError(
+                f"Language {lang} not supported for Repository"
+            )
+        loader = GenericLoader.from_filesystem(
+            path,
+            glob=glob,
+            suffixes=suffixes,
+            exclude=self.exclude_paths,
+            parser=parser,
+            show_progress=True
+        )
+        docs = loader.load()
+        for doc in docs:
+            doc.metadata['url'] = ''
+            doc.metadata['source_type'] = self._source_type
+            doc.metadata['language'] = lang
+        documents = splitter.split_documents(docs)
+        return documents
+
+    def parse(self, source):
+        raise NotImplementedError("Parser method is not implemented for PDFLoader.")