#!/usr/bin/env python
"""Ingest multimodal knowledge base files into structured JSONL."""
from __future__ import annotations
import argparse
import csv
import hashlib
import json
import os
import re
import shutil
import subprocess
import tempfile
from collections import Counter
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Tuple
TEXT_EXTS = {".pdf", ".docx", ".pptx", ".txt", ".md", ".html", ".htm"}
IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".tif", ".tiff", ".bmp"}
AUDIO_EXTS = {".mp3", ".wav", ".m4a", ".flac", ".aac", ".ogg"}
VIDEO_EXTS = {".mp4", ".mov", ".mkv", ".avi", ".webm"}
def normalize_text(text: str) -> str:
text = text.replace("\u0000", "").replace("\ufffd", "")
text = text.replace("\r\n", "\n").replace("\r", "\n")
lines = [re.sub(r"[ \t]+", " ", line).strip() for line in text.split("\n")]
cleaned = []
blank_run = 0
for line in lines:
if not line:
blank_run += 1
if blank_run <= 1:
cleaned.append("")
else:
blank_run = 0
cleaned.append(line)
return "\n".join(cleaned).strip()
def slugify(value: str) -> str:
value = value.strip().lower()
value = re.sub(r"[^a-z0-9]+", "-", value)
value = value.strip("-")
return value or "doc"
def build_doc_id(path: Path, root: Path) -> str:
try:
rel_path = path.relative_to(root).as_posix()
except ValueError:
rel_path = path.as_posix()
slug = slugify(rel_path)
digest = hashlib.sha1(rel_path.encode("utf-8")).hexdigest()[:8]
return f"{slug}__{digest}"
def format_timestamp(seconds: float) -> str:
millis = int(round(seconds * 1000))
hours = millis // 3_600_000
minutes = (millis % 3_600_000) // 60_000
secs = (millis % 60_000) // 1000
ms = millis % 1000
return f"{hours:02d}:{minutes:02d}:{secs:02d}.{ms:03d}"
def safe_int(value: object) -> Optional[int]:
if value is None:
return None
try:
return int(value)
except (TypeError, ValueError):
return None
def detect_file_type(path: Path) -> Optional[str]:
ext = path.suffix.lower()
if ext in TEXT_EXTS:
return ext.lstrip(".")
if ext in IMAGE_EXTS:
return "image"
if ext in AUDIO_EXTS:
return "audio"
if ext in VIDEO_EXTS:
return "video"
if ext == ".csv":
return "csv"
return None
def iter_files(root: Path, max_files: Optional[int] = None) -> Iterable[Path]:
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 max_files is not None and count >= max_files:
break
count += 1
yield path
def load_langchain_documents(path: Path, filetype: str, pdf_strategy: str) -> Tuple[List[object], str]:
try:
from langchain_community.document_loaders import (
TextLoader,
UnstructuredPDFLoader,
UnstructuredWordDocumentLoader,
UnstructuredPowerPointLoader,
UnstructuredHTMLLoader,
UnstructuredMarkdownLoader,
)
except Exception as exc: # pragma: no cover - import guard
raise RuntimeError(
"Missing langchain-community loaders. "
"Install with `pip install -r 01_data_ingestion/requirements.txt`."
) from exc
loader = None
loader_name = ""
if filetype == "pdf":
loader = UnstructuredPDFLoader(str(path), mode="elements", strategy=pdf_strategy)
loader_name = "UnstructuredPDFLoader"
elif filetype == "docx":
loader = UnstructuredWordDocumentLoader(str(path), mode="elements")
loader_name = "UnstructuredWordDocumentLoader"
elif filetype == "pptx":
loader = UnstructuredPowerPointLoader(str(path), mode="elements")
loader_name = "UnstructuredPowerPointLoader"
elif filetype in {"html", "htm"}:
loader = UnstructuredHTMLLoader(str(path), mode="elements")
loader_name = "UnstructuredHTMLLoader"
elif filetype == "md":
loader = UnstructuredMarkdownLoader(str(path), mode="elements")
loader_name = "UnstructuredMarkdownLoader"
elif filetype == "txt":
loader = TextLoader(str(path), encoding="utf-8")
loader_name = "TextLoader"
else:
raise ValueError(f"Unsupported text filetype: {filetype}")
return loader.load(), loader_name
def extract_page(metadata: Dict[str, object]) -> Tuple[Optional[int], Optional[str]]:
for key in ("page_number", "page", "page_num", "page_index", "page_no"):
if key in metadata:
page_val = safe_int(metadata.get(key))
if page_val is not None:
return page_val, "page"
for key in ("slide_number", "slide", "slide_index"):
if key in metadata:
slide_val = safe_int(metadata.get(key))
if slide_val is not None:
return slide_val, "slide"
return None, None
def build_source_base(
path: Path,
root: Path,
filetype: str,
modality: str,
loader_name: str,
) -> Dict[str, object]:
try:
rel_path = path.relative_to(root).as_posix()
except ValueError:
rel_path = path.as_posix()
return {
"path": rel_path,
"filename": path.name,
"filetype": filetype,
"modality": modality,
"loader": loader_name,
}
def add_relations(records: List[Dict[str, object]]) -> None:
total = len(records)
for idx, record in enumerate(records):
record_id = f"{record['doc_id']}:{idx:06d}"
source = record.setdefault("source", {})
source["record_id"] = record_id
source["prev_record_id"] = f"{record['doc_id']}:{idx - 1:06d}" if idx > 0 else None
source["next_record_id"] = (
f"{record['doc_id']}:{idx + 1:06d}" if idx + 1 < total else None
)
def records_from_documents(
docs: List[object],
doc_id: str,
path: Path,
root: Path,
filetype: str,
modality: str,
loader_name: str,
) -> List[Dict[str, object]]:
records: List[Dict[str, object]] = []
for doc in docs:
content = normalize_text(getattr(doc, "page_content", "") or "")
if not content:
continue
metadata = getattr(doc, "metadata", {}) or {}
page_num, page_type = extract_page(metadata)
element_type = metadata.get("category") or metadata.get("element_type")
element_id = metadata.get("element_id") or metadata.get("id")
section_label = None
if page_num is not None:
prefix = "slide" if page_type == "slide" else "page"
section_label = f"{prefix}_{page_num}"
else:
section_label = f"element_{len(records)}"
source = build_source_base(path, root, filetype, modality, loader_name)
if page_num is not None:
source["page_number"] = page_num
source["page_type"] = page_type
if element_type:
source["element_type"] = element_type
if element_id:
source["element_id"] = element_id
records.append(
{
"doc_id": doc_id,
"page": page_num if page_type in {"page", "slide"} else None,
"section": section_label,
"section_index": len(records),
"content": content,
"source": source,
}
)
return records
class OcrRunner:
def __init__(self, engine: str, language: str):
self.engine = engine
self.language = language
self._paddle = None
def extract_text(self, path: Path) -> str:
if self.engine == "tesseract":
try:
from PIL import Image
import pytesseract
except Exception as exc: # pragma: no cover - import guard
raise RuntimeError(
"Missing OCR deps. Install with `pip install -r 01_data_ingestion/requirements.txt` "
"and ensure tesseract is installed."
) from exc
image = Image.open(path)
return pytesseract.image_to_string(image, lang=self.language)
if self.engine == "paddle":
try:
from paddleocr import PaddleOCR
except Exception as exc: # pragma: no cover - import guard
raise RuntimeError(
"Missing paddleocr. Install with `pip install paddleocr`."
) from exc
if self._paddle is None:
lang = "en" if self.language in {"eng", "en"} else self.language
self._paddle = PaddleOCR(use_angle_cls=True, lang=lang, show_log=False)
result = self._paddle.ocr(str(path), cls=True)
lines = []
for page in result:
for item in page:
text = item[1][0]
if text:
lines.append(text)
return "\n".join(lines)
raise ValueError(f"Unsupported OCR engine: {self.engine}")
class WhisperTranscriber:
def __init__(self, model_name: str, device: str, compute_type: str):
try:
from faster_whisper import WhisperModel
except Exception as exc: # pragma: no cover - import guard
raise RuntimeError(
"Missing faster-whisper. Install with `pip install faster-whisper`."
) from exc
self.model = WhisperModel(model_name, device=device, compute_type=compute_type)
def transcribe(self, audio_path: Path):
return self.model.transcribe(
str(audio_path),
beam_size=5,
vad_filter=True,
)
def extract_audio_from_video(video_path: Path, tmp_dir: Path) -> Path:
if shutil.which("ffmpeg") is None:
raise RuntimeError("ffmpeg is required to process video files.")
output_path = tmp_dir / f"{video_path.stem}_audio.wav"
cmd = [
"ffmpeg",
"-y",
"-i",
str(video_path),
"-vn",
"-acodec",
"pcm_s16le",
"-ar",
"16000",
"-ac",
"1",
str(output_path),
]
subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
return output_path
# ---------------------------------------------------------------------------
# Chunked video processing helpers
# ---------------------------------------------------------------------------
def get_video_duration(video_path: Path) -> float:
"""Probe video duration in seconds using ffprobe."""
if shutil.which("ffprobe") is None:
raise RuntimeError("ffprobe is required (usually bundled with ffmpeg).")
cmd = [
"ffprobe", "-v", "error",
"-show_entries", "format=duration",
"-of", "csv=p=0",
str(video_path),
]
result = subprocess.run(cmd, capture_output=True, text=True, check=True)
return float(result.stdout.strip())
def extract_audio_chunk(
video_path: Path,
tmp_dir: Path,
start_sec: float,
duration_sec: float,
chunk_index: int,
) -> Path:
"""Extract a time-bounded audio chunk from a video using FFmpeg."""
if shutil.which("ffmpeg") is None:
raise RuntimeError("ffmpeg is required to process video files.")
output_path = tmp_dir / f"{video_path.stem}_chunk_{chunk_index:04d}.wav"
cmd = [
"ffmpeg", "-y",
"-ss", str(start_sec),
"-i", str(video_path),
"-t", str(duration_sec),
"-vn", "-acodec", "pcm_s16le", "-ar", "16000", "-ac", "1",
str(output_path),
]
subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
return output_path
def load_checkpoint(checkpoint_path: Path) -> Dict:
"""Load a video processing checkpoint, or return empty state."""
if checkpoint_path.exists():
with open(checkpoint_path, "r", encoding="utf-8") as f:
return json.load(f)
return {"completed_chunks": [], "total_segments": 0}
def save_checkpoint(checkpoint_path: Path, state: Dict) -> None:
"""Persist checkpoint state to disk."""
with open(checkpoint_path, "w", encoding="utf-8") as f:
json.dump(state, f, ensure_ascii=True)
# ---------------------------------------------------------------------------
# Keyframe extraction
# ---------------------------------------------------------------------------
def extract_keyframes_ffmpeg(
video_path: Path,
frames_dir: Path,
scene_threshold: float = 0.3,
) -> List[Tuple[Path, float]]:
"""Extract scene-change keyframes from a video using FFmpeg.
Returns list of (frame_path, timestamp_seconds) tuples.
"""
if shutil.which("ffmpeg") is None:
raise RuntimeError("ffmpeg is required for keyframe extraction.")
frames_dir.mkdir(parents=True, exist_ok=True)
cmd = [
"ffmpeg", "-y",
"-i", str(video_path),
"-vf", f"select='gt(scene,{scene_threshold})',showinfo",
"-vsync", "vfr",
str(frames_dir / f"{video_path.stem}_frame_%04d.png"),
]
result = subprocess.run(cmd, capture_output=True, text=True)
# Parse pts_time from showinfo output on stderr
frame_files = sorted(frames_dir.glob(f"{video_path.stem}_frame_*.png"))
timestamps: List[float] = []
for line in result.stderr.split("\n"):
if "pts_time:" in line:
match = re.search(r"pts_time:\s*([\d.]+)", line)
if match:
timestamps.append(float(match.group(1)))
results: List[Tuple[Path, float]] = []
for i, frame_path in enumerate(frame_files):
ts = timestamps[i] if i < len(timestamps) else 0.0
results.append((frame_path, ts))
return results
def extract_keyframes_with_ocr(
video_path: Path,
doc_id: str,
root: Path,
output_dir: Path,
ocr_runner: OcrRunner,
scene_threshold: float = 0.3,
) -> List[Dict[str, object]]:
"""Extract keyframes from a video, run OCR, and produce frame records.
Frame images are saved to <output_dir>/frames/<doc_id>/.
Only frames with OCR text produce JSONL records.
"""
frames_output_dir = output_dir / "frames" / doc_id
keyframes = extract_keyframes_ffmpeg(video_path, frames_output_dir, scene_threshold)
if not keyframes:
return []
try:
rel_video_path = video_path.relative_to(root).as_posix()
except ValueError:
rel_video_path = video_path.as_posix()
records: List[Dict[str, object]] = []
for idx, (frame_path, timestamp_sec) in enumerate(keyframes):
ocr_text = normalize_text(ocr_runner.extract_text(frame_path))
if not ocr_text:
continue
try:
rel_frame_path = frame_path.relative_to(output_dir).as_posix()
except ValueError:
rel_frame_path = frame_path.as_posix()
source: Dict[str, object] = {
"path": rel_video_path,
"filename": video_path.name,
"filetype": "video",
"modality": "video_frame",
"loader": f"ffmpeg+ocr:{ocr_runner.engine}",
"timestamp_start": timestamp_sec,
"timestamp_end": timestamp_sec,
"timestamp": format_timestamp(timestamp_sec),
"frame_path": rel_frame_path,
"frame_index": idx,
"scene_threshold": scene_threshold,
"ocr_engine": ocr_runner.engine,
"ocr_language": ocr_runner.language,
}
records.append({
"doc_id": doc_id,
"page": None,
"section": f"frame_{len(records):04d}",
"section_index": len(records),
"content": ocr_text,
"source": source,
})
return records
# ---------------------------------------------------------------------------
# Chunked video orchestrator
# ---------------------------------------------------------------------------
def process_video_chunked(
video_path: Path,
doc_id: str,
root: Path,
transcriber: WhisperTranscriber,
output_dir: Path,
chunk_duration: float = 600.0,
extract_keyframes: bool = False,
ocr_runner: Optional[OcrRunner] = None,
scene_threshold: float = 0.3,
) -> List[Dict[str, object]]:
"""Process a video in time-based chunks with checkpointing.
Returns all transcript records (and optionally keyframe records)
with globally correct timestamps and segment numbering.
"""
duration = get_video_duration(video_path)
num_chunks = max(1, int(duration // chunk_duration) + (1 if duration % chunk_duration > 0 else 0))
checkpoint_path = output_dir / f".checkpoint_{doc_id}.json"
state = load_checkpoint(checkpoint_path)
completed = set(state["completed_chunks"])
segment_offset = state["total_segments"]
all_transcript_records: List[Dict[str, object]] = []
for chunk_idx in range(num_chunks):
if chunk_idx in completed:
continue
start_sec = chunk_idx * chunk_duration
this_duration = min(chunk_duration, duration - start_sec)
with tempfile.TemporaryDirectory() as tmp_dir:
audio_path = extract_audio_chunk(
video_path, Path(tmp_dir), start_sec, this_duration, chunk_idx
)
chunk_records = records_from_audio(
audio_path, doc_id, root, transcriber, "video",
time_offset=start_sec, segment_offset=segment_offset,
)
# Patch source metadata to point at original video
for rec in chunk_records:
try:
rec["source"]["path"] = video_path.relative_to(root).as_posix()
except ValueError:
rec["source"]["path"] = video_path.as_posix()
rec["source"]["filename"] = video_path.name
rec["source"]["filetype"] = "video"
all_transcript_records.extend(chunk_records)
segment_offset += len(chunk_records)
completed.add(chunk_idx)
state["completed_chunks"] = sorted(completed)
state["total_segments"] = segment_offset
save_checkpoint(checkpoint_path, state)
# Keyframe extraction (optional)
frame_records: List[Dict[str, object]] = []
if extract_keyframes and ocr_runner is not None:
frame_records = extract_keyframes_with_ocr(
video_path, doc_id, root, output_dir, ocr_runner,
scene_threshold=scene_threshold,
)
# Link transcript and frame records independently
if all_transcript_records:
add_relations(all_transcript_records)
if frame_records:
add_relations(frame_records)
# Clean up checkpoint on success
if checkpoint_path.exists():
checkpoint_path.unlink()
return all_transcript_records + frame_records
def records_from_csv(path: Path, doc_id: str, root: Path) -> List[Dict[str, object]]:
records: List[Dict[str, object]] = []
with open(path, newline="", encoding="utf-8") as handle:
reader = csv.DictReader(handle)
for row_idx, row in enumerate(reader, start=1):
parts = []
for key, value in row.items():
if value is None:
continue
field = (key or "").strip()
if not field:
field = "column"
cleaned = str(value).strip()
if not cleaned:
continue
parts.append(f"{field}: {cleaned}")
content = normalize_text(" | ".join(parts))
if not content:
continue
source = build_source_base(path, root, "csv", "table", "csv")
source["row_index"] = row_idx
source["column_count"] = len(reader.fieldnames or [])
records.append(
{
"doc_id": doc_id,
"page": None,
"section": f"row_{row_idx}",
"section_index": len(records),
"content": content,
"source": source,
}
)
return records
def records_from_image(
path: Path,
doc_id: str,
root: Path,
ocr_runner: OcrRunner,
) -> List[Dict[str, object]]:
content = normalize_text(ocr_runner.extract_text(path))
if not content:
return []
source = build_source_base(path, root, "image", "image", f"ocr:{ocr_runner.engine}")
source["ocr_engine"] = ocr_runner.engine
source["ocr_language"] = ocr_runner.language
return [
{
"doc_id": doc_id,
"page": None,
"section": "image_1",
"section_index": 0,
"content": content,
"source": source,
}
]
def records_from_audio(
path: Path,
doc_id: str,
root: Path,
transcriber: WhisperTranscriber,
modality: str,
time_offset: float = 0.0,
segment_offset: int = 0,
) -> List[Dict[str, object]]:
segments, info = transcriber.transcribe(path)
language = getattr(info, "language", None)
language_prob = getattr(info, "language_probability", None)
records: List[Dict[str, object]] = []
for idx, segment in enumerate(segments):
text = normalize_text(getattr(segment, "text", "") or "")
if not text:
continue
start = float(getattr(segment, "start", 0.0)) + time_offset
end = float(getattr(segment, "end", 0.0)) + time_offset
global_idx = segment_offset + len(records)
source = build_source_base(path, root, modality, modality, "faster-whisper")
source["timestamp_start"] = start
source["timestamp_end"] = end
source["timestamp"] = format_timestamp(start)
if language:
source["language"] = language
if language_prob is not None:
source["language_probability"] = round(float(language_prob), 4)
records.append(
{
"doc_id": doc_id,
"page": None,
"section": f"segment_{global_idx:04d}",
"section_index": global_idx,
"content": text,
"source": source,
}
)
return records
def _resolve_modality(filetype: str) -> str:
"""Map a filetype string to its modality category."""
if filetype == "image":
return "image"
if filetype == "audio":
return "audio"
if filetype == "video":
return "video"
if filetype == "csv":
return "table"
return "text"
def ingest(
input_dir: Path,
output_dir: Path,
output_name: str,
errors_name: str,
pdf_strategy: str,
ocr_engine: str,
ocr_language: str,
whisper_model: str,
whisper_device: str,
whisper_compute_type: str,
max_files: Optional[int],
chunk_duration: float = 600.0,
extract_keyframes: bool = False,
scene_threshold: float = 0.3,
) -> Dict[str, object]:
output_dir.mkdir(parents=True, exist_ok=True)
output_path = output_dir / output_name
errors_path = output_dir / errors_name
file_counts = Counter()
record_counts = 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]] = []
skipped_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,
modality: str = "", filetype: str = "") -> None:
nonlocal error_count
error_count += 1
error_entry: Dict[str, object] = {
"path": str(path),
"stage": stage,
"error": str(err),
}
if modality:
error_entry["modality"] = modality
if filetype:
error_entry["filetype"] = filetype
error_handle.write(json.dumps(error_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(input_dir, max_files=max_files):
filetype = detect_file_type(path)
if not filetype:
skipped_files.append({
"path": _rel(path),
"reason": f"Unsupported file type: {path.suffix}",
})
continue
total_files += 1
file_counts[filetype] += 1
modality = _resolve_modality(filetype)
doc_id = build_doc_id(path, input_dir)
try:
records: List[Dict[str, object]] = []
if filetype == "csv":
records = records_from_csv(path, doc_id, input_dir)
elif filetype == "image":
records = records_from_image(path, doc_id, input_dir, ocr_runner)
elif filetype == "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 filetype == "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, filetype, pdf_strategy)
records = records_from_documents(
docs, doc_id, path, input_dir, filetype, "text", loader_name
)
if records:
# Video records already have relations from process_video_chunked
if filetype != "video":
add_relations(records)
for record in records:
out_handle.write(json.dumps(record, ensure_ascii=True) + "\n")
total_records += len(records)
record_counts[filetype] += len(records)
success_files.append(_rel(path))
else:
empty_files.append({
"path": _rel(path),
"modality": modality,
"reason": "Extraction produced no content",
})
log_error(path, "empty_content",
Exception("Extraction produced no content"),
modality=modality, filetype=filetype)
except Exception as exc:
log_error(path, "ingest", exc,
modality=modality, filetype=filetype)
failed_files.append({
"path": _rel(path),
"modality": modality,
"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": {
"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),
"files_skipped": len(skipped_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,
"skipped_files": skipped_files,
}
stats_path = output_dir / "ingestion_stats.json"
with open(stats_path, "w", encoding="utf-8") as stats_handle:
json.dump(stats, stats_handle, indent=2, ensure_ascii=True)
return stats
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Ingest multimodal knowledge base files.")
parser.add_argument(
"--input",
default=os.path.join("00_Data"),
help="Root directory containing files to ingest.",
)
parser.add_argument(
"--output-dir",
default=os.path.join("01_data_ingestion", "output"),
help="Output directory for JSONL and stats.",
)
parser.add_argument(
"--output-name",
default="cleaned_documents.jsonl",
help="Output JSONL filename.",
)
parser.add_argument(
"--errors-name",
default="ingestion_errors.jsonl",
help="Errors JSONL filename.",
)
parser.add_argument(
"--pdf-strategy",
default="auto",
help="Unstructured PDF strategy (auto, hi_res, fast, ocr_only).",
)
parser.add_argument(
"--ocr-engine",
default="tesseract",
choices=["tesseract", "paddle"],
help="OCR engine for images.",
)
parser.add_argument(
"--ocr-language",
default="eng",
help="OCR language code (tesseract or paddle).",
)
parser.add_argument(
"--whisper-model",
default="small",
help="faster-whisper model size (tiny, base, small, medium, large-v3).",
)
parser.add_argument(
"--whisper-device",
default="cpu",
help="faster-whisper device (cpu, cuda, auto).",
)
parser.add_argument(
"--whisper-compute-type",
default="int8",
help="faster-whisper compute type (int8, float16, int8_float16).",
)
parser.add_argument(
"--max-files",
type=int,
default=None,
help="Optional cap on number of files processed.",
)
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).",
)
parser.add_argument(
"--workers",
type=int,
default=4,
help="Number of parallel worker threads (0 = sequential legacy mode, default: 4).",
)
return parser.parse_args()
def main() -> None:
args = parse_args()
if args.workers > 0:
from ingestion.engine import IngestionConfig, IngestionEngine
from ingestion.events import LogEmitter
config = IngestionConfig(
input_dir=Path(args.input),
output_dir=Path(args.output_dir),
output_name=args.output_name,
errors_name=args.errors_name,
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,
workers=args.workers,
)
engine = IngestionEngine(config, emitter=LogEmitter())
engine.run()
else:
ingest(
input_dir=Path(args.input),
output_dir=Path(args.output_dir),
output_name=args.output_name,
errors_name=args.errors_name,
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,
)
if __name__ == "__main__":
main()