68 lines
2.1 KiB
Python
Executable File
68 lines
2.1 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
if str(PROJECT_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(PROJECT_ROOT))
|
|
|
|
from src.rag.extractors.v16_1_2.v16_extractor import V16Extractor
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Run the formal V16.1.2 plant-science SRT extraction pipeline."
|
|
)
|
|
parser.add_argument("--pdf", required=True, type=Path)
|
|
parser.add_argument("--output-dir", required=True, type=Path)
|
|
parser.add_argument("--pmid", default="")
|
|
parser.add_argument("--doi", default="")
|
|
parser.add_argument("--disable-llm", action="store_true")
|
|
parser.add_argument("--use-pc-data", action="store_true")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
pdf = args.pdf.resolve()
|
|
output_dir = args.output_dir.resolve()
|
|
if not pdf.exists():
|
|
raise SystemExit(f"PDF does not exist: {pdf}")
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
extractor = V16Extractor(
|
|
output_dir=output_dir,
|
|
use_mineru=True,
|
|
use_llm_relations=not args.disable_llm,
|
|
use_rule_relation_fallback=False,
|
|
use_vlm=False,
|
|
use_neo4j=False,
|
|
use_unified_storage=False,
|
|
use_llm_definitions=False,
|
|
use_pc_data=args.use_pc_data,
|
|
use_disambiguation=args.use_pc_data,
|
|
)
|
|
result = extractor.extract(
|
|
pdf,
|
|
pmid=args.pmid or None,
|
|
doi=args.doi or None,
|
|
)
|
|
summary = {
|
|
"version": result.get("metadata", {}).get("version"),
|
|
"pdf": str(pdf),
|
|
"output_dir": str(output_dir),
|
|
"entities": len(result.get("entities", [])),
|
|
"relations": len(result.get("relations", [])),
|
|
"relation_units": len(result.get("chunks", [])),
|
|
"result_json": str(output_dir / "result.json"),
|
|
}
|
|
print(json.dumps(summary, ensure_ascii=False, indent=2))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|