Spaces:
Running on Zero
Running on Zero
| """Small command-line entry point; paths are relative to the working directory.""" | |
| import json | |
| from pathlib import Path | |
| from typing import Annotated | |
| import typer | |
| from visual_doc_rag import data | |
| app = typer.Typer(no_args_is_help=True, pretty_exceptions_enable=False) | |
| Subset = Annotated[str, typer.Option(help="Dataset subset: hr or industrial.")] | |
| DataDir = Annotated[Path, typer.Option(help="Local data/cache directory; ignored by Git.")] | |
| def check_subset(subset: str) -> None: | |
| if subset not in data.SUBSETS: | |
| raise typer.BadParameter(f"Choose one of: {', '.join(data.SUBSETS)}", param_hint="--subset") | |
| def build_demo( | |
| subset: Subset = "hr", | |
| data_dir: DataDir = Path("data"), | |
| output: Annotated[Path | None, typer.Option(help="New bundle directory.")] = None, | |
| model: Annotated[ | |
| str, typer.Option(help="Model used by the existing sample cache.") | |
| ] = "gemini-3.5-flash-lite", | |
| space: Annotated[bool, typer.Option(help="Include Hugging Face Space configuration.")] = False, | |
| ) -> None: | |
| """Package offline retrieval and five cached samples, without calling the provider.""" | |
| from visual_doc_rag.demo import build_demo as pack | |
| check_subset(subset) | |
| try: | |
| result = pack(subset, data_dir, output, model, space) | |
| except Exception as error: | |
| typer.echo(f"Demo build failed: {error}", err=True) | |
| raise typer.Exit(1) from error | |
| typer.echo(f"Demo bundle: {result}") | |
| def download( | |
| subset: Subset = "hr", | |
| data_dir: DataDir = Path("data"), | |
| revision: Annotated[ | |
| str, typer.Option(help="HF branch, tag or commit; resolved to a SHA.") | |
| ] = "main", | |
| ) -> None: | |
| """Download corpus, queries, qrels and document metadata (no original PDFs).""" | |
| check_subset(subset) | |
| try: | |
| manifest = data.download(subset, data_dir, revision) | |
| except Exception as error: | |
| typer.echo(f"Download failed: {error}", err=True) | |
| raise typer.Exit(1) from error | |
| typer.echo(json.dumps(manifest, indent=2)) | |
| def inspect(subset: Subset = "hr", data_dir: DataDir = Path("data")) -> None: | |
| """Inspect local tables, validate joins and save aggregate statistics/sample pages.""" | |
| check_subset(subset) | |
| try: | |
| report = data.inspect(subset, data_dir) | |
| except Exception as error: | |
| typer.echo(f"Inspection failed: {error}", err=True) | |
| raise typer.Exit(1) from error | |
| typer.echo(json.dumps({k: v for k, v in report.items() if k != "features"}, indent=2)) | |
| typer.echo(f"Saved schema and statistics: {data_dir / subset / 'inspection.json'}") | |
| if not all(report["integrity_checks"].values()): | |
| typer.echo("Integrity checks failed; review inspection.json before indexing.", err=True) | |
| raise typer.Exit(1) | |
| def pages(subset: Subset = "hr", data_dir: DataDir = Path("data")) -> None: | |
| """Export all original page images to a revision-aware PNG cache.""" | |
| from visual_doc_rag.pages import export_pages | |
| check_subset(subset) | |
| try: | |
| result = export_pages(subset, data_dir) | |
| except Exception as error: | |
| typer.echo(f"Page export failed: {error}", err=True) | |
| raise typer.Exit(1) from error | |
| typer.echo(f"Exported {len(result['pages'])} pages; {result['image_bytes'] / 1e6:.1f} MB") | |
| def index( | |
| subset: Subset = "hr", | |
| data_dir: DataDir = Path("data"), | |
| device: Annotated[str, typer.Option(help="auto, cuda or cpu for page encoding.")] = "auto", | |
| mode: Annotated[str, typer.Option(help="all, text or visual.")] = "all", | |
| model_revision: Annotated[ | |
| str, typer.Option(help="Adapter revision; reuse cached pin by default.") | |
| ] = "main", | |
| base_revision: Annotated[ | |
| str | None, typer.Option(help="Optional pinned base-model revision.") | |
| ] = None, | |
| ) -> None: | |
| """Build BM25 and visual indexes; reuse completed page embeddings on retry.""" | |
| import torch | |
| from visual_doc_rag.text_index import TextIndex | |
| from visual_doc_rag.visual_index import build_visual | |
| check_subset(subset) | |
| if mode not in {"all", "text", "visual"} or device not in {"auto", "cuda", "cpu"}: | |
| raise typer.BadParameter("Use --mode all/text/visual and --device auto/cuda/cpu.") | |
| torch.set_num_threads(4) | |
| if device == "auto": | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| try: | |
| source, tables = data.load_tables(subset, data_dir) | |
| if mode in {"all", "text"}: | |
| result = TextIndex.build(tables["corpus"], data_dir / subset / "indexes", source) | |
| typer.echo(f"BM25: {json.dumps(result)}") | |
| if mode in {"all", "visual"}: | |
| result = build_visual(subset, data_dir, source, device, model_revision, base_revision) | |
| typer.echo(json.dumps({k: v for k, v in result.items() if k != "pages"}, indent=2)) | |
| except Exception as error: | |
| typer.echo(f"Index build failed: {error}", err=True) | |
| raise typer.Exit(1) from error | |
| def eval_command( | |
| subset: Subset = "hr", | |
| data_dir: DataDir = Path("data"), | |
| runs_dir: Annotated[ | |
| Path, typer.Option(help="Directory for timestamped metric reports.") | |
| ] = Path("runs"), | |
| ) -> None: | |
| """Evaluate text/visual/hybrid on every English query, with CPU latency.""" | |
| import torch | |
| from visual_doc_rag.evaluate import evaluate, format_table | |
| check_subset(subset) | |
| torch.set_num_threads(4) | |
| try: | |
| output, report = evaluate(subset, data_dir, runs_dir) | |
| except Exception as error: | |
| typer.echo(f"Evaluation failed: {error}", err=True) | |
| raise typer.Exit(1) from error | |
| typer.echo(format_table(report)) | |
| typer.echo(f"Queries: {report['n_queries']}; pages: {report['n_pages']}") | |
| typer.echo(f"Saved: {output / 'metrics.json'}") | |
| def search( | |
| query: Annotated[str, typer.Argument(help="Question or keywords to retrieve pages for.")], | |
| subset: Subset = "hr", | |
| data_dir: DataDir = Path("data"), | |
| method: Annotated[str, typer.Option(help="text, visual or hybrid.")] = "hybrid", | |
| top_k: Annotated[int, typer.Option(min=1, help="Number of page results to display.")] = 5, | |
| export: Annotated[ | |
| bool, typer.Option(help="Save an HTML report, page images and JSON under runs/search/.") | |
| ] = False, | |
| ) -> None: | |
| """Retrieve pages and print local image paths, without generating an answer.""" | |
| from visual_doc_rag.retrieve import PageRetriever | |
| check_subset(subset) | |
| if method not in {"text", "visual", "hybrid"} or not query.strip(): | |
| raise typer.BadParameter("Provide a nonempty query and --method text/visual/hybrid.") | |
| try: | |
| retriever = PageRetriever(subset, data_dir) | |
| results = retriever.search(query, method, top_k) | |
| output = {"method": method, "pages": results} | |
| if export: | |
| from visual_doc_rag.search_export import export_search | |
| output["exported_report"] = str( | |
| export_search(query, method, retriever.source, results, retriever.tables["corpus"]) | |
| ) | |
| except Exception as error: | |
| typer.echo(f"Search failed: {error}", err=True) | |
| raise typer.Exit(1) from error | |
| typer.echo(json.dumps(output, indent=2)) | |
| def ask( | |
| query: Annotated[str, typer.Argument(help="Question to answer using retrieved page images.")], | |
| subset: Subset = "hr", | |
| data_dir: DataDir = Path("data"), | |
| method: Annotated[str, typer.Option(help="text, visual or hybrid.")] = "hybrid", | |
| export: Annotated[bool, typer.Option(help="Save an answer and evidence HTML report.")] = False, | |
| ) -> None: | |
| """Answer from the top-3 pages with validated citations (requires GEMINI_API_KEY).""" | |
| from visual_doc_rag.serve import AnswerService | |
| check_subset(subset) | |
| if method not in {"text", "visual", "hybrid"} or not query.strip(): | |
| raise typer.BadParameter("Provide a nonempty query and --method text/visual/hybrid.") | |
| try: | |
| service = AnswerService(subset, data_dir) | |
| output = service.ask(query, method) | |
| if export: | |
| from visual_doc_rag.search_export import export_search | |
| output["exported_report"] = str( | |
| export_search( | |
| output["query"], | |
| method, | |
| service.retriever.source, | |
| output["pages"], | |
| service.retriever.tables["corpus"], | |
| answer=output, | |
| ) | |
| ) | |
| except Exception as error: | |
| typer.echo(f"Answer failed: {error}", err=True) | |
| raise typer.Exit(1) from error | |
| typer.echo(json.dumps(output, indent=2, ensure_ascii=False)) | |