#!/usr/bin/env python
"""Batch ingest all files of a specific modality from a folder.
Wraps the processing functions from ingest_multimodal.py and adds
modality/filetype filtering, progress reporting, and dependency pre-checks.
Usage:
# All text files in a folder
python ingest_folder.py --input /path/to/docs --modality text --output-dir output/
# Only PDFs
python ingest_folder.py --input /path/to/docs --filetype pdf --output-dir output/
# All images with PaddleOCR
python ingest_folder.py --input /path/to/images --modality image --ocr-engine paddle
# Everything in a folder (no filter)
python ingest_folder.py --input /path/to/data --output-dir output/
"""
from __future__ import annotations
import argparse
import json
import shutil
import sys
from collections import Counter
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, Iterable, List, Optional
# Ensure sibling module is importable
sys.path.insert(0, str(Path(__file__).resolve().parent))
from ingest_multimodal import (
TEXT_EXTS,
IMAGE_EXTS,
AUDIO_EXTS,
VIDEO_EXTS,
detect_file_type,
build_doc_id,
add_relations,
load_langchain_documents,
records_from_documents,
records_from_image,
records_from_audio,
records_from_csv,
process_video_chunked,
OcrRunner,
WhisperTranscriber,
_resolve_modality,
)
# ---------------------------------------------------------------------------
# Modality mapping
# ---------------------------------------------------------------------------
MODALITY_EXTS = {
"text": TEXT_EXTS,
"image": IMAGE_EXTS,
"audio": AUDIO_EXTS,
"video": VIDEO_EXTS,
"table": {".csv"},
}
# ---------------------------------------------------------------------------
# Filtered file iterator
# ---------------------------------------------------------------------------
def iter_files_filtered(
root: Path,
modality: Optional[str] = None,
filetype: Optional[str] = None,
max_files: Optional[int] = None,
) -> Iterable[Path]:
"""Recursively discover files with optional modality/filetype filtering.
Args:
root: Root directory to search.
modality: One of 'text', 'image', 'audio', 'video', 'table', or None for all.
filetype: Specific extension (without dot) like 'pdf', 'png', 'flac'.
Overrides modality if both are provided.
max_files: Optional cap on number of files yielded.
"""
allowed_exts = None
if filetype:
ext = f".{filetype.lower().lstrip('.')}"
allowed_exts = {ext}
elif modality and modality != "all":
allowed_exts = MODALITY_EXTS.get(modality)
if allowed_exts is None:
raise ValueError(
f"Unknown modality '{modality}'. Choose from: {', '.join(MODALITY_EXTS)}"
)
count = 0
for path in sorted(root.rglob("*")):
if path.is_dir():
continue
if path.name.startswith("."):
continue
if "output" in path.parts:
continue
if allowed_exts is not None:
if path.suffix.lower() not in allowed_exts:
continue
if detect_file_type(path) is None:
continue
if max_files is not None and count >= max_files:
break
count += 1
yield path
# ---------------------------------------------------------------------------
# Dependency pre-checks
# ---------------------------------------------------------------------------
def check_dependencies(modality: Optional[str], filetype: Optional[str]) -> None:
"""Verify required system tools and Python packages are available."""
needs_ocr = False
needs_whisper = False
needs_ffmpeg = False
if filetype:
ext = f".{filetype.lower().lstrip('.')}"
if ext in IMAGE_EXTS:
needs_ocr = True
elif ext in AUDIO_EXTS:
needs_whisper = True
elif ext in VIDEO_EXTS:
needs_whisper = True
needs_ffmpeg = True
elif modality == "image":
needs_ocr = True
elif modality == "audio":
needs_whisper = True
elif modality == "video":
needs_whisper = True
needs_ffmpeg = True
elif modality in (None, "all"):
# Can't pre-check for all, dependencies checked per-file at runtime
return
errors = []
if needs_ocr:
if shutil.which("tesseract") is None:
try:
import pytesseract # noqa: F401
except ImportError:
errors.append("tesseract is not installed. Install with: brew install tesseract")
if needs_whisper:
try:
import faster_whisper # noqa: F401
except ImportError:
errors.append(
"faster-whisper is not installed. "
"Install with: pip install faster-whisper (requires Python <3.13)"
)
if needs_ffmpeg:
if shutil.which("ffmpeg") is None:
errors.append("ffmpeg is not installed. Install with: brew install ffmpeg")
if shutil.which("ffprobe") is None:
errors.append("ffprobe is not installed (usually bundled with ffmpeg).")
if errors:
print("Missing dependencies:")
for err in errors:
print(f" - {err}")
sys.exit(1)
# ---------------------------------------------------------------------------
# Main ingestion
# ---------------------------------------------------------------------------
def ingest_folder(
input_dir: Path,
output_dir: Path,
output_name: str = "cleaned_documents.jsonl",
errors_name: str = "ingestion_errors.jsonl",
modality: Optional[str] = None,
filetype: Optional[str] = None,
pdf_strategy: str = "auto",
ocr_engine: str = "tesseract",
ocr_language: str = "eng",
whisper_model: str = "small",
whisper_device: str = "cpu",
whisper_compute_type: str = "int8",
max_files: Optional[int] = None,
chunk_duration: float = 600.0,
extract_keyframes: bool = False,
scene_threshold: float = 0.3,
) -> Dict:
"""Ingest all matching files from a folder into downstream-ready JSONL."""
output_dir.mkdir(parents=True, exist_ok=True)
output_path = output_dir / output_name
errors_path = output_dir / errors_name
file_counts: Counter = Counter()
record_counts: Counter = Counter()
error_count = 0
total_files = 0
total_records = 0
success_files: List[str] = []
failed_files: List[Dict[str, str]] = []
empty_files: List[Dict[str, str]] = []
ocr_runner = OcrRunner(ocr_engine, ocr_language)
transcriber = None
def _rel(path: Path) -> str:
try:
return path.relative_to(input_dir).as_posix()
except ValueError:
return path.as_posix()
def log_error(path: Path, stage: str, err: Exception,
mod: str = "", ft: str = "") -> None:
nonlocal error_count
error_count += 1
entry: Dict[str, object] = {
"path": str(path), "stage": stage, "error": str(err),
}
if mod:
entry["modality"] = mod
if ft:
entry["filetype"] = ft
error_handle.write(json.dumps(entry, ensure_ascii=True) + "\n")
with open(output_path, "w", encoding="utf-8") as out_handle, \
open(errors_path, "w", encoding="utf-8") as error_handle:
for path in iter_files_filtered(input_dir, modality, filetype, max_files):
ft = detect_file_type(path)
if not ft:
continue
total_files += 1
file_counts[ft] += 1
mod = _resolve_modality(ft)
doc_id = build_doc_id(path, input_dir)
# Progress reporting
if total_files % 100 == 0:
print(f" [{total_files}] Processing: {path.name}")
try:
records: List[Dict] = []
if ft == "csv":
records = records_from_csv(path, doc_id, input_dir)
elif ft == "image":
records = records_from_image(path, doc_id, input_dir, ocr_runner)
elif ft == "audio":
if transcriber is None:
transcriber = WhisperTranscriber(
whisper_model, whisper_device, whisper_compute_type
)
records = records_from_audio(
path, doc_id, input_dir, transcriber, "audio"
)
elif ft == "video":
if transcriber is None:
transcriber = WhisperTranscriber(
whisper_model, whisper_device, whisper_compute_type
)
records = process_video_chunked(
video_path=path,
doc_id=doc_id,
root=input_dir,
transcriber=transcriber,
output_dir=output_dir,
chunk_duration=chunk_duration,
extract_keyframes=extract_keyframes,
ocr_runner=ocr_runner if extract_keyframes else None,
scene_threshold=scene_threshold,
)
else:
docs, loader_name = load_langchain_documents(path, ft, pdf_strategy)
records = records_from_documents(
docs, doc_id, path, input_dir, ft, "text", loader_name
)
if records:
# Video records already have relations from process_video_chunked
if ft != "video":
add_relations(records)
for rec in records:
out_handle.write(json.dumps(rec, ensure_ascii=True) + "\n")
total_records += len(records)
record_counts[ft] += len(records)
success_files.append(_rel(path))
else:
empty_files.append({
"path": _rel(path),
"modality": mod,
"reason": "Extraction produced no content",
})
log_error(path, "empty_content",
Exception("Extraction produced no content"),
mod=mod, ft=ft)
except Exception as exc:
log_error(path, "ingest", exc, mod=mod, ft=ft)
failed_files.append({
"path": _rel(path),
"modality": mod,
"reason": str(exc),
})
stats = {
"input_root": str(input_dir),
"output_file": str(output_path),
"errors_file": str(errors_path),
"generated_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
"config": {
"modality_filter": modality or "all",
"filetype_filter": filetype,
"pdf_strategy": pdf_strategy,
"ocr_engine": ocr_engine,
"ocr_language": ocr_language,
"whisper_model": whisper_model,
"whisper_device": whisper_device,
"whisper_compute_type": whisper_compute_type,
"video_chunk_duration": chunk_duration,
"extract_keyframes": extract_keyframes,
"scene_threshold": scene_threshold if extract_keyframes else None,
},
"counts": {
"files_processed": total_files,
"files_succeeded": len(success_files),
"files_failed": len(failed_files),
"files_empty": len(empty_files),
"records_emitted": total_records,
"errors": error_count,
},
"files_by_type": dict(file_counts),
"records_by_type": dict(record_counts),
"success_files": success_files,
"failed_files": failed_files,
"empty_files": empty_files,
}
stats_path = output_dir / "ingestion_stats.json"
with open(stats_path, "w", encoding="utf-8") as f:
json.dump(stats, f, indent=2, ensure_ascii=True)
return stats
# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Batch ingest all files of a specific type from a folder.",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
python ingest_folder.py --input /path/to/docs --modality text
python ingest_folder.py --input /path/to/docs --filetype pdf
python ingest_folder.py --input /path/to/images --modality image --ocr-engine paddle
python ingest_folder.py --input /path/to/data --output-dir output/
""",
)
parser.add_argument("--input", required=True, help="Root directory containing files to ingest.")
parser.add_argument(
"--output-dir",
default="01_data_ingestion/output",
help="Output directory for JSONL and stats (default: 01_data_ingestion/output).",
)
parser.add_argument("--output-name", default="cleaned_documents.jsonl", help="Output JSONL filename.")
parser.add_argument(
"--modality",
choices=["text", "image", "audio", "video", "table", "all"],
default="all",
help="Filter by modality (default: all).",
)
parser.add_argument(
"--filetype",
default=None,
help="Filter by specific extension (e.g. pdf, png, flac). Overrides --modality.",
)
parser.add_argument("--max-files", type=int, default=None, help="Optional cap on files processed.")
parser.add_argument("--pdf-strategy", default="auto", help="Unstructured PDF strategy.")
parser.add_argument("--ocr-engine", default="tesseract", choices=["tesseract", "paddle"])
parser.add_argument("--ocr-language", default="eng", help="OCR language code.")
parser.add_argument("--whisper-model", default="small", help="faster-whisper model size.")
parser.add_argument("--whisper-device", default="cpu", help="faster-whisper device.")
parser.add_argument("--whisper-compute-type", default="int8", help="faster-whisper compute type.")
parser.add_argument("--chunk-duration", type=float, default=600.0,
help="Duration in seconds for video audio chunks (default: 600 = 10 min).")
parser.add_argument("--extract-keyframes", action="store_true", default=False,
help="Extract keyframes from videos and run OCR on them.")
parser.add_argument("--scene-threshold", type=float, default=0.3,
help="FFmpeg scene change threshold 0.0-1.0 (default: 0.3).")
return parser.parse_args()
def main() -> None:
args = parse_args()
input_dir = Path(args.input)
if not input_dir.exists():
print(f"Error: input directory '{input_dir}' not found.")
sys.exit(1)
effective_modality = None if args.filetype else (args.modality if args.modality != "all" else None)
check_dependencies(effective_modality, args.filetype)
print(f"Ingesting from: {input_dir}")
if args.filetype:
print(f" Filter: filetype={args.filetype}")
elif args.modality != "all":
print(f" Filter: modality={args.modality}")
else:
print(" Filter: none (all supported types)")
print(f" Output: {args.output_dir}\n")
stats = ingest_folder(
input_dir=input_dir,
output_dir=Path(args.output_dir),
output_name=args.output_name,
modality=effective_modality,
filetype=args.filetype,
pdf_strategy=args.pdf_strategy,
ocr_engine=args.ocr_engine,
ocr_language=args.ocr_language,
whisper_model=args.whisper_model,
whisper_device=args.whisper_device,
whisper_compute_type=args.whisper_compute_type,
max_files=args.max_files,
chunk_duration=args.chunk_duration,
extract_keyframes=args.extract_keyframes,
scene_threshold=args.scene_threshold,
)
counts = stats["counts"]
print(f"\nIngestion complete!")
print(f" Files processed: {counts['files_processed']}")
print(f" Succeeded: {counts['files_succeeded']}")
print(f" Failed: {counts['files_failed']}")
print(f" Empty (no text): {counts['files_empty']}")
print(f" Records emitted: {counts['records_emitted']}")
print(f" Files by type: {stats['files_by_type']}")
print(f" Records by type: {stats['records_by_type']}")
print(f" Output: {stats['output_file']}")
if counts["errors"] > 0:
print(f" Errors log: {stats['errors_file']}")
if __name__ == "__main__":
main()