73 lines
2.9 KiB
Python
73 lines
2.9 KiB
Python
 
|
#!/usr/bin/env python3
|
||
|
|
"""CLI to chunk documents and print/save results."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import json
|
||
|
|
import sys
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "backend"))
|
||
|
|
|
||
|
|
from rag_cut.models import SplitConfig, SplitMode
|
||
|
|
from rag_cut.pipeline import chunk_document
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
parser = argparse.ArgumentParser(description="RAG-cut document chunking CLI")
|
||
|
|
parser.add_argument("file", type=Path, help="Path to document")
|
||
|
|
parser.add_argument(
|
||
|
|
"--mode",
|
||
|
|
choices=[m.value for m in SplitMode],
|
||
|
|
default="default",
|
||
|
|
help="Split mode",
|
||
|
|
)
|
||
|
|
parser.add_argument("--delimiter", default=None, help="Delimiter for delimiter mode")
|
||
|
|
parser.add_argument("--parent-delimiter", default=None, help="Parent delimiter for parent_child mode")
|
||
|
|
parser.add_argument("--child-delimiter", default=None, help="Child delimiter for parent_child mode")
|
||
|
|
parser.add_argument("--max-chunk-size", type=int, default=1500, help="Parent/primary max chunk size")
|
||
|
|
parser.add_argument("--child-max-size", type=int, default=512, help="Child max size for parent_child mode")
|
||
|
|
parser.add_argument("--overlap", type=int, default=150)
|
||
|
|
parser.add_argument("--header-row-start", type=int, default=1)
|
||
|
|
parser.add_argument("--header-row-end", type=int, default=1)
|
||
|
|
parser.add_argument("--start-row", type=int, default=2)
|
||
|
|
parser.add_argument("--rows-per-chunk", type=int, default=1)
|
||
|
|
parser.add_argument("-o", "--output", type=Path, default=None, help="Save JSON result")
|
||
|
|
parser.add_argument("--preview", type=int, default=3, help="Print first N chunks")
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
config = SplitConfig(
|
||
|
|
mode=SplitMode(args.mode),
|
||
|
|
delimiter=args.delimiter,
|
||
|
|
parent_delimiter=args.parent_delimiter,
|
||
|
|
child_delimiter=args.child_delimiter,
|
||
|
|
max_chunk_size=args.max_chunk_size,
|
||
|
|
child_max_size=args.child_max_size,
|
||
|
|
overlap=args.overlap,
|
||
|
|
header_row_start=args.header_row_start,
|
||
|
|
header_row_end=args.header_row_end,
|
||
|
|
start_row=args.start_row,
|
||
|
|
rows_per_chunk=args.rows_per_chunk,
|
||
|
|
)
|
||
|
|
|
||
|
|
result = chunk_document(args.file, config=config)
|
||
|
|
print(f"File: {result.filename}")
|
||
|
|
print(f"Doc ID: {result.doc_id}")
|
||
|
|
print(f"Blocks: {result.block_count} -> Chunks: {result.chunk_count}")
|
||
|
|
print(f"Mode: {result.split_mode.value}")
|
||
|
|
print("-" * 60)
|
||
|
|
|
||
|
|
for chunk in result.chunks[: args.preview]:
|
||
|
|
print(f"\n[Chunk {chunk.index}] ({chunk.char_count} chars) types={chunk.block_types}")
|
||
|
|
preview = chunk.content[:500]
|
||
|
|
print(preview + ("..." if len(chunk.content) > 500 else ""))
|
||
|
|
|
||
|
|
if args.output:
|
||
|
|
args.output.write_text(result.model_dump_json(indent=2), encoding="utf-8")
|
||
|
|
print(f"\nSaved full result to {args.output}")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|