ai-parrot 0.1.0__cp311-cp311-manylinux_2_28_x86_64.whl

parrot/loaders/image.py

@@ -0,0 +1,119 @@
+from typing import Any
+from collections.abc import Callable
+from pathlib import Path, PurePath
+import numpy as np
+from PIL import Image
+from langchain.docstore.document import Document
+from transformers import CLIPModel
+import torch
+from torchvision import transforms
+from .abstract import AbstractLoader
+from ..stores.abstract import AbstractStore
+
+
+class ImageLoader(AbstractLoader):
+    """
+    Image Loader.
+    """
+    _extension = ['.jpg', '.jpeg', '.png']
+    chunk_size = 768
+
+    def __init__(
+        self,
+        path: PurePath,
+        store: AbstractStore,
+        tokenizer: Callable[..., Any] = None,
+        text_splitter: Callable[..., Any] = None,
+        source_type: str = 'image',
+        **kwargs
+    ):
+        super().__init__(tokenizer, text_splitter, source_type, **kwargs)
+        self.path = path
+        if isinstance(path, str):
+            self.path = Path(path).resolve()
+        # Model:
+        self._model = CLIPModel.from_pretrained(
+            # "openai/clip-vit-base-patch32"
+            "openai/clip-vit-large-patch14-336"
+        )
+        # Define image preprocessing
+        self._preprocess = transforms.Compose(
+            [
+                transforms.Resize((336, 336)),  # Adjust the size to match the model's expected input
+                transforms.CenterCrop(336),  # Optionally add a center crop if needed
+                transforms.ToTensor(),
+                transforms.Normalize(
+                    (0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)
+                )  # CLIP's original normalization
+            ]
+        )
+        # required Milvus Store:
+        self.store = store
+
+    def transform_image(self, img_data):
+        image = self._preprocess(img_data)
+        image = image.unsqueeze(0)
+        with torch.no_grad():
+            features = self._model.get_image_features(pixel_values=image)
+        embedding = features.squeeze().cpu().numpy()
+        return embedding.astype(np.float32)
+
+    def _insert_image(self, data):
+        return self.store.insert(data)
+
+    def _load_image(self, path) -> list:
+        """
+        Load an Image file.
+        Args:
+            path (Path): The path to the Image file.
+        Returns:
+            list: A list of Langchain Documents.
+        """
+        if self._check_path(path):
+            self.logger.info(f"Loading Image file: {path}")
+            img = Image.open(path).convert('RGB')
+            embedding = self.transform_image(img).tolist()
+            data={
+                "url": '',
+                "source": f"{path.name}",
+                "filename": path,
+                "question": '',
+                "answer": '',
+                "source_type": self._source_type,
+                "type": "image",
+                "text": '',
+                "vector": embedding,
+                "document_meta": {
+                    "image": path.name,
+                    "extension": path.suffix
+                }
+            }
+            self._insert_image([embedding])
+        return []
+
+    def load(self) -> list:
+        """
+        Load data from a Image file.
+        Returns:
+            list: A list of Langchain Documents.
+        """
+        if not self.path.exists():
+            raise FileNotFoundError(f"Image file/directory not found: {self.path}")
+        if self.path.is_dir():
+            # iterate over the files in the directory
+            for ext in self._extension:
+                for item in self.path.glob(f'*{ext}'):
+                    self._load_image(item)
+        elif self.path.is_file():
+            self._load_image(self.path)
+        else:
+            raise ValueError(
+                f"Image Loader: Invalid path: {self.path}"
+            )
+        # Load Image loads the image directly to database.
+        return True
+
+    def parse(self, source):
+        raise NotImplementedError(
+            "Parser method is not implemented for ImageLoader."
+        )

parrot/loaders/json.py

@@ -0,0 +1,52 @@
+from collections.abc import Callable
+from pathlib import PurePath
+from langchain_community.document_loaders import JSONLoader as JSLoader
+from .abstract import AbstractLoader
+
+
+class JSONLoader(AbstractLoader):
+    """
+    Loader for JSON files.
+    """
+    _extension = ['.json']
+    extract_metadata: Callable = None
+
+    def extract_metadata(self, record: dict, metadata: dict) -> dict:
+        meta = {
+            "source_type": self._source_type,
+            "priority": self._priority,
+        }
+        return meta
+
+    def load(self, path: PurePath) -> list:
+        """
+        Load data from a JSON file.
+
+        Args:
+            source (str): The path to the JSON file.
+
+        Returns:
+            list: A list of Langchain Documents.
+        """
+        if self._check_path(path):
+            self.logger.info(f"Loading JSON file: {path}")
+            # Create metadata for each chunk
+            meta = {
+                "filename": str(path),
+            }
+            args = {
+                "metadata_func": self.extract_metadata,
+            }
+            loader = JSLoader(
+                file_path=path,
+                jq_schema=".",
+                text_content=False,
+                **args
+            )
+            documents = loader.load()
+            for doc in documents:
+                doc.metadata.update(meta)
+            # Split the documents into chunks
+            return self.split_documents(documents)
+        else:
+            return []

parrot/loaders/pdf.py

@@ -0,0 +1,187 @@
+from collections.abc import Callable
+from pathlib import Path, PurePath
+from typing import Any
+from io import BytesIO
+import fitz
+import pytesseract
+from PIL import Image
+from langchain.docstore.document import Document
+from .basepdf import BasePDF
+
+
+class PDFLoader(BasePDF):
+    """
+    Loader for 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
+        )
+        self.parse_images = kwargs.get('parse_images', False)
+        # Table Settings:
+        self.table_settings = {
+            #"vertical_strategy": "text",
+            # "horizontal_strategy": "text",
+            "intersection_x_tolerance": 3,
+            "intersection_y_tolerance": 3
+        }
+        table_settings = kwargs.get('table_setttings', {})
+        if table_settings:
+            self.table_settings.update(table_settings)
+
+    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]
+                text = page.get_text()
+                # first: text
+                if text:
+                    page_num = page_number + 1
+                    try:
+                        summary = self.get_summary_from_text(text)
+                    except Exception:
+                        summary = ''
+                    metadata = {
+                        "url": '',
+                        "source": f"{path.name} Page.#{page_num}",
+                        "filename": path.name,
+                        "index": f"{page_num}",
+                        "type": 'pdf',
+                        "question": '',
+                        "answer": '',
+                        "source_type": self._source_type,
+                        "data": {},
+                        "summary": summary,
+                        "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", ""),
+                        }
+                    }
+                    docs.append(
+                        Document(
+                            page_content=text,
+                            metadata=metadata
+                        )
+                    )
+                # Extract images and use OCR to get text from each image
+                # second: images
+                if self.parse_images is True:
+                    image_list = page.get_images(full=True)
+                    file_name = path.stem.replace(' ', '_').replace('.', '').lower()
+                    for img_index, img in enumerate(image_list):
+                        xref = img[0]
+                        base_image = pdf.extract_image(xref)
+                        image = Image.open(BytesIO(base_image["image"]))
+                        url = ''
+                        if self.save_images is True:
+                            img_name = f'image_{file_name}_{page_num}_{img_index}.png'
+                            img_path = self._imgdir.joinpath(img_name)
+                            self.logger.notice(
+                                f"Saving Image Page on {img_path}"
+                            )
+                            try:
+                                image.save(
+                                    img_path,
+                                    format="png",
+                                    optimize=True
+                                )
+                                url = f'/static/images/{img_name}'
+                            except OSError:
+                                pass
+                        # Use Tesseract to extract text from image
+                        image_text = pytesseract.image_to_string(
+                            image,
+                            lang=self._lang
+                        )
+                        # TODO: add the summary (explanation)
+                        # Create a document for each image
+                        image_meta = {
+                            "url": url,
+                            "source": f"{path.name} Page.#{page_num}",
+                            "filename": path.name,
+                            "index": f"{path.name}:{page_num}",
+                            "question": '',
+                            "answer": '',
+                            "type": 'image',
+                            "data": {},
+                            "summary": '',
+                            "document_meta": {
+                                "image_index": img_index,
+                                "image_name": img_name,
+                                "description": f"Extracted from {page_number}."
+                            },
+                            "source_type": self._source_type
+                        }
+                        docs.append(
+                            Document(page_content=image_text, metadata=image_meta)
+                        )
+                # third: tables
+                # Look for tables on this page and display the table count
+                try:
+                    tabs = page.find_tables()
+                    for tab_idx, tab in enumerate(tabs):
+                        # iterating over all tables in page:
+                        df = tab.to_pandas()  # convert to pandas DataFrame
+                        # converting to markdown, but after pre-processing pandas
+                        df = df.dropna(axis=1, how='all')
+                        df = df.dropna(how='all', axis=0)  # Drop empty rows
+                        table_meta = {
+                            "url": '',
+                            "source": f"{path.name} Page.#{page_num} Table.#{tab_idx}",
+                            "filename": path.name,
+                            "index": f"{path.name}:{page_num}",
+                            "question": '',
+                            "answer": '',
+                            "type": 'table',
+                            "data": {},
+                            "summary": '',
+                            "document_meta": {
+                                "table_index": tab_idx,
+                                "table_shape": df.shape,
+                                "table_columns": df.columns.tolist(),
+                                "description": f"Extracted from {page_number}."
+                            },
+                            "source_type": self._source_type
+                        }
+                        txt = df.to_markdown()
+                        if txt:
+                            docs.append(
+                                Document(page_content=txt, metadata=table_meta)
+                            )
+                except Exception as exc:
+                    print(exc)
+            pdf.close()
+            return docs
+        else:
+            return []

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 []