diff --git a/utils/pdf_processing.py b/utils/pdf_processing.py index 4ec3a46..a8893f0 100644 --- a/utils/pdf_processing.py +++ b/utils/pdf_processing.py @@ -1,25 +1,17 @@ -import asyncio import gc -from io import BytesIO - import pdfplumber - +from io import BytesIO +import asyncio from utils.chunking import split_text_to_subchunks -from utils.config import MAX_PROCESSES_DEEPSEEK, MAX_PROCESSES_GROQ -from utils.regular_helpers import elements_to_positions, extract_page_content -from utils.scanned_helpers import ( - deepseek_translate_worker, - is_scanned_page, - process_scanned_page_worker, -) - +from utils.config import MAX_PROCESSES_DEEPSEEK, MAX_PROCESSES_GROQ, PDF_BATCH_SIZE +from utils.regular_helpers import extract_page_content, elements_to_positions +from utils.scanned_helpers import is_scanned_page, process_scanned_page_worker, deepseek_translate_worker async def groq_worker(job, semaphore): async with semaphore: loop = asyncio.get_running_loop() return await loop.run_in_executor(None, process_scanned_page_worker, job) - - + async def deepseek_worker(job, semaphore): async with semaphore: loop = asyncio.get_running_loop() @@ -30,74 +22,43 @@ async def deepseek_worker(job, semaphore): gc.collect() return sub_chunks - -async def process_pdf(pdf_stream): - pdf_bytes = pdf_stream.read() +async def process_pdf_batch(pdf_bytes, start_page=0, end_page=None): all_sub_chunks = [] scanned_jobs = [] - total_pages = 0 with pdfplumber.open(BytesIO(pdf_bytes)) as pdf: total_pages = len(pdf.pages) - print(f"Total pages: {total_pages}") + if end_page is None or end_page > total_pages: + end_page = total_pages - for i, page in enumerate(pdf.pages): + for i in range(start_page, end_page): + page = pdf.pages[i] scanned = is_scanned_page(page) - print(f"\nProcessing Page {i + 1}: {'SCANNED' if scanned else 'TEXT'}") - if scanned: - continue - # scanned_jobs.append((i, pdf_bytes)) + scanned_jobs.append((i, pdf_bytes)) else: elements = extract_page_content(page) positions = elements_to_positions(elements) for pos in positions: sub_chunks = split_text_to_subchunks( - pos["content"], - i + 1, - pos["position"], - pos["type"], - is_scanned=False, + pos["content"], i+1, pos["position"], pos["type"], is_scanned=False ) all_sub_chunks.extend(sub_chunks) - scanned_pages_count = len(scanned_jobs) - regular_pages_count = total_pages - scanned_pages_count - groq_results = [] - if scanned_jobs: - print( - f"\nProcessing {scanned_pages_count} scanned pages with Groq in batches of {MAX_PROCESSES_GROQ}" - ) groq_semaphore = asyncio.Semaphore(MAX_PROCESSES_GROQ) groq_tasks = [groq_worker(job, groq_semaphore) for job in scanned_jobs] groq_results = await asyncio.gather(*groq_tasks) deepseek_results = [] - if groq_results: deepseek_jobs = [(res["page"], res["raw_content"]) for res in groq_results] - print( - f"\nTranslating {len(deepseek_jobs)} pages with DeepSeek using {MAX_PROCESSES_DEEPSEEK} processes" - ) deepseek_semaphore = asyncio.Semaphore(MAX_PROCESSES_DEEPSEEK) - deepseek_tasks = [ - deepseek_worker(job, deepseek_semaphore) for job in deepseek_jobs - ] + deepseek_tasks = [deepseek_worker(job, deepseek_semaphore) for job in deepseek_jobs] deepseek_results = await asyncio.gather(*deepseek_tasks) for sub_chunks in deepseek_results: all_sub_chunks.extend(sub_chunks) - del pdf_bytes, scanned_jobs - if "groq_results" in locals(): - del groq_results - if "deepseek_results" in locals(): - del deepseek_results gc.collect() - - return { - "chunks": all_sub_chunks, - "scanned_pages": scanned_pages_count, - "regular_pages": regular_pages_count, - } + return all_sub_chunks, len(scanned_jobs), (end_page - start_page - len(scanned_jobs))