私有RAG之极限优化等待时间

                ┌───────────────┐
                │   Web Server  │
                │   (Flask)     │
                └───────┬───────┘
                        │
                        │查询
                        ▼
                ┌───────────────┐
                │  Vector DB    │
                └───────────────┘


                文档处理流水线（后台）
私有 API
      │
      ▼
Document Queue
      │
      ▼
Parser Workers
      │
      ▼
Embedding Workers
      │
      ▼
Vector DB

import aiohttp
import asyncio

async def download_attachment(session, mail):

    url = mail["url"]
    filename = f"pdfs/{mail['id']}.pdf"

    async with session.get(url) as resp:
        data = await resp.read()

    with open(filename, "wb") as f:
        f.write(data)

    return filename


async def download_all(mails):

    async with aiohttp.ClientSession() as session:

        tasks = [
            download_attachment(session, mail)
            for mail in mails
        ]

        results = await asyncio.gather(*tasks)

    return results


asyncio.run(download_all(mails))

私有API
↓
下载
↓
Amazon S3
↓
解析

def has_text(page):
    text = page.get_text()
    return len(text.strip()) > 60

def parse_page(page):
    return page.get_text()

def parse_pdf(file_path):
    with fitz.open(file_path) as doc:
        with ThreadPoolExecutor(max_workers=8) as executor:
    		executor.map(pipeline, pdf_files)

...
for r in executor.map(...):
    process(r)   		

PyTorch/Paddle
    ↓
ONNX
    ↓
ONNX Runtime

规则 chunk（最快）
semantic chunk（embedding）
layout chunk（复杂文档）

chunk generator
      │
      ▼
chunk queue
      │
      ▼
Worker Pool (2~4)
      │
      ▼
Batch Builder (64 chunks)
      │
      ▼
ONNX Runtime 推理
      │
      ▼
Vector DB Batch Insert

import onnxruntime as ort
from concurrent.futures import ProcessPoolExecutor

session = ort.InferenceSession(
    "embedding.onnx",
    providers=["CPUExecutionProvider"]
)

def embed_batch(texts):

    inputs = tokenizer(
        texts,
        padding=True,
        truncation=True,
        return_tensors="np"
    )

    outputs = session.run(None, dict(inputs))

    return outputs[0]

def process_chunks(chunk_batch):

    embeddings = embed_batch(chunk_batch)

    vector_db.insert(embeddings)

def run_pipeline(chunks):

    batches = list(batch(chunks,64))

    with ProcessPoolExecutor(max_workers=4) as executor:
        executor.map(process_chunks, batches)

load_embedding_model()
load_ocr_model()