fix: support content type for input data

Author: ChenZiHong-Gavin

graphgen/models/reader/csv_reader.py

@@ -9,6 +9,9 @@ class CSVReader(BaseReader):
     def read(self, file_path: str) -> List[Dict[str, Any]]:
 
         df = pd.read_csv(file_path)
-        if self.text_column not in df.columns:
-            raise ValueError(f"Missing '{self.text_column}' column in CSV file.")
-        return df.to_dict(orient="records")
+        for _, row in df.iterrows():
+            if row.get("type") == "text" and self.text_column not in row:
+                raise ValueError(
+                    f"Missing '{self.text_column}' in document: {row.to_dict()}"
+                )
+        return self.filter(df.to_dict(orient="records"))

graphgen/models/reader/json_reader.py

@@ -10,9 +10,9 @@ def read(self, file_path: str) -> List[Dict[str, Any]]:
             data = json.load(f)
             if isinstance(data, list):
                 for doc in data:
-                    if self.text_column not in doc:
+                    if doc["type"] == "text" and self.text_column not in doc:
                         raise ValueError(
                             f"Missing '{self.text_column}' in document: {doc}"
                         )
-                return data
+                return self.filter(data)
             raise ValueError("JSON file must contain a list of documents.")

graphgen/models/reader/jsonl_reader.py

@@ -12,12 +12,11 @@ def read(self, file_path: str) -> List[Dict[str, Any]]:
             for line in f:
                 try:
                     doc = json.loads(line)
-                    if self.text_column in doc:
-                        docs.append(doc)
-                    else:
+                    if doc["type"] == "text" and self.text_column not in doc:
                         raise ValueError(
                             f"Missing '{self.text_column}' in document: {doc}"
                         )
+                    docs.append(doc)
                 except json.JSONDecodeError as e:
                     logger.error("Error decoding JSON line: %s. Error: %s", line, e)
-        return docs
+        return self.filter(docs)

graphgen/models/reader/pdf_reader.py

@@ -74,7 +74,7 @@ def read(self, file_path: str, **override) -> List[Dict[str, Any]]:
         kwargs = {**self._default_kwargs, **override}
 
         mineru_result = self._call_mineru(pdf_path, kwargs)
-        return mineru_result
+        return self.filter(mineru_result)
 
     def _call_mineru(
         self, pdf_path: Path, kwargs: Dict[str, Any]
@@ -172,8 +172,6 @@ def _try_load_cached_result(
             for key in ("page_idx", "bbox", "text_level"):
                 if item.get(key) is not None:
                     del item[key]
-            if item["type"] == "text" and not item["content"].strip():
-                continue
             results.append(item)
         return results
 

graphgen/models/reader/txt_reader.py

@@ -11,4 +11,4 @@ def read(self, file_path: str) -> List[Dict[str, Any]]:
                 line = line.strip()
                 if line:
                     docs.append({self.text_column: line})
-        return docs
+        return self.filter(docs)