Add image_min_size parameter

backend/rag_1/config.py

@@ -4,7 +4,7 @@
 
 from omegaconf import OmegaConf
 from omegaconf.dictconfig import DictConfig
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, validator
 from pydantic.dataclasses import dataclass
 
 
@@ -41,9 +41,20 @@ class IngestConfig:
 
     metadata_keys: list[str]
     table_format: Literal["text", "html", "image"]
+    image_min_size: list[float]
+    table_min_size: list[float]
 
     export_extracted: bool
 
+    @validator("image_min_size", "table_min_size")
+    def validate_size(cls, value: list[float]) -> list[float]:
+        """Check that the value is between 0 and 1."""
+        if len(value) != 2:
+            raise ValueError("Size must be a list of two floats.")
+        if min(value) < 0 or max(value) > 1:
+            raise ValueError("Size must be a list of floats between 0 and 1.")
+        return value
+
 
 @dataclass(config=ConfigDict(extra="forbid"))
 class Config:

backend/rag_1/config.yaml

@@ -48,5 +48,7 @@ ingest:
     - filename
     - page_number
   table_format: "html" # "text" or "html" or "image"
+  image_min_size: [0.1, 0.1]
+  table_min_size: [0.0, 0.0]
 
   export_extracted: True

backend/rag_1/ingest.py

@@ -39,7 +39,11 @@ def ingest_pdf(file_path: str | Path, config: DictConfig) -> None:
     )
 
     # Get images
-    images = select_images(raw_pdf_elements, config.ingest.metadata_keys)
+    images = select_images(
+        raw_pdf_elements,
+        metadata_keys=config.ingest.metadata_keys,
+        min_size=config.ingest.image_min_size,
+    )
 
     # Get chunks
     if config.ingest.chunking_enable:
@@ -49,9 +53,12 @@ def ingest_pdf(file_path: str | Path, config: DictConfig) -> None:
         chunks = raw_pdf_elements
 
     # Get text, tables
-    texts = select_texts(chunks, config.ingest.metadata_keys)
+    texts = select_texts(chunks, metadata_keys=config.ingest.metadata_keys)
     tables = select_tables(
-        chunks, config.ingest.table_format, config.ingest.metadata_keys
+        chunks,
+        table_format=config.ingest.table_format,
+        metadata_keys=config.ingest.metadata_keys,
+        min_size=config.ingest.table_min_size,
     )
 
     vectorstore = get_vectorstore(config)

backend/rag_1/notebook.ipynb

@@ -192,7 +192,11 @@
    "outputs": [],
    "source": [
     "# Get images\n",
-    "images = select_images(raw_pdf_elements)\n",
+    "images = select_images(\n",
+    "    raw_pdf_elements,\n",
+    "    metadata_keys=config.ingest.metadata_keys,\n",
+    "    min_size=config.ingest.image_min_size,\n",
+    ")\n",
     "for image in images[:N_DISPLAY]:\n",
     "    display(image)"
    ]
@@ -233,8 +237,13 @@
    "outputs": [],
    "source": [
     "# Get text, tables\n",
-    "texts = select_texts(chunks, config.ingest.metadata_keys)\n",
-    "tables = select_tables(chunks, config.ingest.table_format, config.ingest.metadata_keys)"
+    "texts = select_texts(chunks, metadata_keys=config.ingest.metadata_keys)\n",
+    "tables = select_tables(\n",
+    "    chunks,\n",
+    "    table_format=config.ingest.table_format,\n",
+    "    metadata_keys=config.ingest.metadata_keys,\n",
+    "    min_size=config.ingest.table_min_size,\n",
+    ")"
    ]
   },
   {

backend/rag_2/config.py

@@ -4,7 +4,7 @@
 
 from omegaconf import OmegaConf
 from omegaconf.dictconfig import DictConfig
-from pydantic import BaseModel, ConfigDict, root_validator
+from pydantic import BaseModel, ConfigDict, root_validator, validator
 from pydantic.dataclasses import dataclass
 
 
@@ -50,6 +50,8 @@ class IngestConfig:
 
     metadata_keys: list[str]
     table_format: Literal["text", "html", "image"]
+    image_min_size: list[float]
+    table_min_size: list[float]
 
     summarize_text: bool
     summarize_table: bool
@@ -118,6 +120,15 @@ def validate_fields(cls, values: dict) -> dict:
 
         return values
 
+    @validator("image_min_size", "table_min_size")
+    def validate_size(cls, value: list[float]) -> list[float]:
+        """Check that the value is between 0 and 1."""
+        if len(value) != 2:
+            raise ValueError("Size must be a list of two floats.")
+        if min(value) < 0 or max(value) > 1:
+            raise ValueError("Size must be a list of floats between 0 and 1.")
+        return value
+
 
 @dataclass(config=ConfigDict(extra="forbid"))
 class Config:

backend/rag_2/config.yaml

@@ -65,6 +65,8 @@ ingest:
     - filename
     - page_number
   table_format: "html" # "text" or "html" or "image"
+  image_min_size: [0.1, 0.1]
+  table_min_size: [0.0, 0.0]
 
   summarize_text: False
   summarize_table: True

backend/rag_2/ingest.py

@@ -150,7 +150,11 @@ async def ingest_pdf(file_path: str | Path, config: DictConfig) -> None:
     )
 
     # Get images
-    images = select_images(raw_pdf_elements, config.ingest.metadata_keys)
+    images = select_images(
+        raw_pdf_elements,
+        metadata_keys=config.ingest.metadata_keys,
+        min_size=config.ingest.image_min_size,
+    )
 
     # Get chunks
     if config.ingest.chunking_enable:
@@ -160,9 +164,12 @@ async def ingest_pdf(file_path: str | Path, config: DictConfig) -> None:
         chunks = raw_pdf_elements
 
     # Get text, tables
-    texts = select_texts(chunks, config.ingest.metadata_keys)
+    texts = select_texts(chunks, metadata_keys=config.ingest.metadata_keys)
     tables = select_tables(
-        chunks, config.ingest.table_format, config.ingest.metadata_keys
+        chunks,
+        table_format=config.ingest.table_format,
+        metadata_keys=config.ingest.metadata_keys,
+        min_size=config.ingest.table_min_size,
     )
 
     # Summarize text

backend/rag_2/notebook.ipynb

@@ -198,7 +198,11 @@
    "outputs": [],
    "source": [
     "# Get images\n",
-    "images = select_images(raw_pdf_elements)\n",
+    "images = select_images(\n",
+    "    raw_pdf_elements,\n",
+    "    metadata_keys=config.ingest.metadata_keys,\n",
+    "    min_size=config.ingest.image_min_size,\n",
+    ")\n",
     "for image in images[:N_DISPLAY]:\n",
     "    display(image)"
    ]
@@ -239,8 +243,13 @@
    "outputs": [],
    "source": [
     "# Get text, tables\n",
-    "texts = select_texts(chunks, config.ingest.metadata_keys)\n",
-    "tables = select_tables(chunks, config.ingest.table_format, config.ingest.metadata_keys)"
+    "texts = select_texts(chunks, metadata_keys=config.ingest.metadata_keys)\n",
+    "tables = select_tables(\n",
+    "    chunks,\n",
+    "    table_format=config.ingest.table_format,\n",
+    "    metadata_keys=config.ingest.metadata_keys,\n",
+    "    min_size=config.ingest.table_min_size,\n",
+    ")"
    ]
   },
   {

backend/rag_3/config.py

@@ -4,7 +4,7 @@
 
 from omegaconf import OmegaConf
 from omegaconf.dictconfig import DictConfig
-from pydantic import BaseModel, ConfigDict, root_validator
+from pydantic import BaseModel, ConfigDict, root_validator, validator
 from pydantic.dataclasses import dataclass
 
 
@@ -44,12 +44,13 @@ class IngestConfig:
     """Configuration for PDF ingestion."""
 
     clear_database: bool
-
     chunking_enable: bool
     chunking_func: HydraObject
 
     metadata_keys: list[str]
     table_format: Literal["text", "html", "image"]
+    image_min_size: list[float]
+    table_min_size: list[float]
 
     summarize_text: bool
     summarize_table: bool
@@ -108,6 +109,15 @@ def validate_fields(cls, values: dict) -> dict:
 
         return values
 
+    @validator("image_min_size", "table_min_size")
+    def validate_size(cls, value: list[float]) -> list[float]:
+        """Check that the value is between 0 and 1."""
+        if len(value) != 2:
+            raise ValueError("Size must be a list of two floats.")
+        if min(value) < 0 or max(value) > 1:
+            raise ValueError("Size must be a list of floats between 0 and 1.")
+        return value
+
 
 @dataclass(config=ConfigDict(extra="forbid"))
 class Config:

backend/rag_3/config.yaml

@@ -65,6 +65,8 @@ ingest:
     - filename
     - page_number
   table_format: "image" # "text" or "html" or "image"
+  image_min_size: [0.1, 0.1]
+  table_min_size: [0.0, 0.0]
 
   summarize_text: False
   summarize_table: True

backend/rag_3/ingest.py

@@ -150,7 +150,11 @@ async def ingest_pdf(file_path: str | Path, config: DictConfig) -> None:
     )
 
     # Get images
-    images = select_images(raw_pdf_elements, config.ingest.metadata_keys)
+    images = select_images(
+        raw_pdf_elements,
+        metadata_keys=config.ingest.metadata_keys,
+        min_size=config.ingest.image_min_size,
+    )
 
     # Get chunks
     if config.ingest.chunking_enable:
@@ -160,9 +164,12 @@ async def ingest_pdf(file_path: str | Path, config: DictConfig) -> None:
         chunks = raw_pdf_elements
 
     # Get text, tables
-    texts = select_texts(chunks, config.ingest.metadata_keys)
+    texts = select_texts(chunks, metadata_keys=config.ingest.metadata_keys)
     tables = select_tables(
-        chunks, config.ingest.table_format, config.ingest.metadata_keys
+        chunks,
+        table_format=config.ingest.table_format,
+        metadata_keys=config.ingest.metadata_keys,
+        min_size=config.ingest.table_min_size,
     )
 
     # Summarize text

backend/rag_3/notebook.ipynb

@@ -198,7 +198,11 @@
    "outputs": [],
    "source": [
     "# Get images\n",
-    "images = select_images(raw_pdf_elements)\n",
+    "images = select_images(\n",
+    "    raw_pdf_elements,\n",
+    "    metadata_keys=config.ingest.metadata_keys,\n",
+    "    min_size=config.ingest.image_min_size,\n",
+    ")\n",
     "for image in images[:N_DISPLAY]:\n",
     "    display(image)"
    ]
@@ -239,8 +243,16 @@
    "outputs": [],
    "source": [
     "# Get text, tables\n",
-    "texts = select_texts(chunks, config.ingest.metadata_keys)\n",
-    "tables = select_tables(chunks, config.ingest.table_format, config.ingest.metadata_keys)"
+    "texts = select_texts(\n",
+    "    chunks,\n",
+    "    metadata_keys=config.ingest.metadata_keys,\n",
+    ")\n",
+    "tables = select_tables(\n",
+    "    chunks,\n",
+    "    table_format=config.ingest.table_format,\n",
+    "    metadata_keys=config.ingest.metadata_keys,\n",
+    "    min_size=config.ingest.table_min_size,\n",
+    ")"
    ]
   },
   {