feat: filter non-exist content

Author: ChenZiHong-Gavin

graphgen/bases/base_reader.py

@@ -1,6 +1,9 @@
+import os
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List
 
+import requests
+
 
 class BaseReader(ABC):
     """
@@ -18,3 +21,45 @@ def read(self, file_path: str) -> List[Dict[str, Any]]:
         :param file_path: Path to the input file.
         :return: List of dictionaries containing the data.
         """
+
+    @staticmethod
+    def filter(data: List[dict]) -> List[dict]:
+        """
+        Filter out entries with empty or missing text in the specified column.
+
+        :param data: List of dictionaries containing the data.
+        :return: Filtered list of dictionaries.
+        """
+
+        def _image_exists(path_or_url: str, timeout: int = 3) -> bool:
+            """
+            Check if an image exists at the given local path or URL.
+            :param path_or_url: Local file path or remote URL of the image.
+            :param timeout: Timeout for remote URL requests in seconds.
+            :return: True if the image exists, False otherwise.
+            """
+            if not path_or_url:
+                return False
+
+            if not path_or_url.startswith(("http://", "https://", "ftp://")):
+                path = path_or_url.replace("file://", "", 1)
+                return os.path.isfile(path)
+            try:
+                resp = requests.head(path_or_url, allow_redirects=True, timeout=timeout)
+                return resp.status_code == 200
+            except requests.RequestException:
+                return False
+
+        filtered_data = []
+        for item in data:
+            if item.get("type") == "text":
+                content = item.get("content", "").strip()
+                if content:
+                    filtered_data.append(item)
+            elif item.get("type") in ("image", "table", "equation"):
+                img_path = item.get("img_path")
+                if _image_exists(img_path):
+                    filtered_data.append(item)
+            else:
+                filtered_data.append(item)
+        return filtered_data

graphgen/configs/vqa_config.yaml

@@ -1,5 +1,5 @@
 read:
-  input_file: resources/input_examples/pdf_demo.pdf # input file path, support json, jsonl, txt, pdf. See resources/input_examples for examples
+  input_file: resources/input_examples/vqa_demo.json # input file path, support json, jsonl, txt, pdf. See resources/input_examples for examples
 split:
   chunk_size: 1024 # chunk size for text splitting
   chunk_overlap: 100 # chunk overlap for text splitting

graphgen/graphgen.py

@@ -92,6 +92,7 @@ async def insert(self, read_config: Dict, split_config: Dict):
         """
         # Step 1: Read files
         data = read_files(read_config["input_file"], self.working_dir)
+        print(data)
         if len(data) == 0:
             logger.warning("No data to process")
             return