Async fast endpoint

huridocs · Jul 19, 2024 · ba16548 · ba16548
1 parent 3cb0648
commit ba16548
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/src/app.py b/src/app.py
@@ -31,14 +31,16 @@ async def error():
 @catch_exceptions
 async def run(file: UploadFile = File(...), fast: bool = Form(False)):
  if fast:
- return await run_in_threadpool(analyze_pdf_fast, file.file.read())
+ return await run_in_threadpool(analyze_pdf_fast, file.file.read(), "")
 
  return await run_in_threadpool(analyze_pdf, file.file.read(), "")
 
 
 @app.post("/save_xml/{xml_file_name}")
 @catch_exceptions
-async def analyze_and_save_xml(file: UploadFile = File(...), xml_file_name: str | None = None):
+async def analyze_and_save_xml(file: UploadFile = File(...), xml_file_name: str | None = None, fast: bool = Form(False)):
+ if fast:
+ return await run_in_threadpool(analyze_pdf_fast, file.file.read(), xml_file_name)
  return await run_in_threadpool(analyze_pdf, file.file.read(), xml_file_name)
 
 

diff --git a/src/pdf_layout_analysis/run_pdf_layout_analysis_fast.py b/src/pdf_layout_analysis/run_pdf_layout_analysis_fast.py
@@ -12,10 +12,10 @@
 from data_model.SegmentBox import SegmentBox
 
 
-def analyze_pdf_fast(file: AnyStr):
+def analyze_pdf_fast(file: AnyStr, xml_file_name: str = "") -> list[dict]:
  pdf_path = pdf_content_to_pdf_path(file)
  service_logger.info("Creating Paragraph Tokens [fast]")
- pdf_features = PdfFeatures.from_pdf_path(pdf_path)
+ pdf_features = PdfFeatures.from_pdf_path(pdf_path, xml_file_name)
  token_type_trainer = TokenTypeTrainer([pdf_features], ModelConfiguration())
  token_type_trainer.set_token_types(join(ROOT_PATH, "models", "token_type_lightgbm.model"))
  trainer = ParagraphExtractorTrainer(pdfs_features=[pdf_features], model_configuration=PARAGRAPH_EXTRACTION_CONFIGURATION)