用视觉嵌入 + Agent 自动生成 IDP 文档 Schema

# 依赖：pip install transformers torch scikit-learn pillow pdf2image
# 还需安装 poppler（pdf2image 的底层）：brew install poppler  (macOS) 或 apt install poppler-utils (Linux)

import os
import numpy as np
from pdf2image import convert_from_path
from PIL import Image
from transformers import AutoModel, AutoProcessor
from sklearn.cluster import AgglomerativeClustering

# 1. 把 PDF 转成图片并提取视觉嵌入
model_name = "google/siglip-so400m-patch14-384"  # 视觉编码器，也可换用其他 ViT 模型
processor = AutoProcessor.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

def pdf_to_images(pdf_path: str, dpi: int = 200) -> list[Image.Image]:
    """将 PDF 每页转为 PIL Image，只取第一页做代表。"""
    images = convert_from_path(pdf_path, dpi=dpi, first_page=1, last_page=1)
    return images[0]

def get_visual_embedding(image: Image.Image) -> np.ndarray:
    """单页文档 -> 视觉嵌入向量。"""
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        embeddings = model.get_image_features(**inputs)
    # 归一化，方便后续用余弦距离聚类
    emb = embeddings.cpu().numpy().flatten()
    emb = emb / np.linalg.norm(emb)
    return emb

# 假设你的文档都在这个目录下
doc_dir = "./raw_docs"
pdf_files = [f for f in os.listdir(doc_dir) if f.lower().endswith(".pdf")]

embeddings = []
for pdf in pdf_files:
    img = pdf_to_images(os.path.join(doc_dir, pdf))
    emb = get_visual_embedding(img)
    embeddings.append(emb)

X = np.stack(embeddings)  # shape: (num_docs, emb_dim)

# 2. 聚类——不预设簇数，用距离阈值自动决定
clusterer = AgglomerativeClustering(
    metric="cosine",      # 视觉嵌入用余弦距离更合理
    linkage="average",
    distance_threshold=0.35,  # 阈值需根据你的文档集微调
    n_clusters=None       # 不固定簇数，由阈值决定
)
labels = clusterer.fit_predict(X)

# 3. 输出聚类结果
for cluster_id in set(labels):
    members = [pdf_files[i] for i in range(len(labels)) if labels[i] == cluster_id]
    print(f"簇 {cluster_id}（{len(members)} 份文档）: {members}")

# 依赖：pip install openai pytesseract pillow pdf2image
# 还需安装 Tesseract OCR：brew install tesseract (macOS) 或 apt install tesseract-ocr (Linux)

import json
import os
import pytesseract
from pdf2image import convert_from_path
from openai import OpenAI

client = OpenAI()  # 默认读取环境变量 OPENAI_API_KEY

def extract_text_from_pdf(pdf_path: str) -> str:
    """OCR 抽取 PDF 第一页文本，供 Agent 阅读。"""
    images = convert_from_path(pdf_path, dpi=200, first_page=1, last_page=1)
    text = pytesseract.image_to_string(images[0], lang="chi_sim+eng")
    return text

def generate_schema(sample_texts: list[str], cluster_label: str) -> dict:
    """让 Agent 根据采样文档文本生成提取 Schema。"""
    combined = "\n---\n".join(sample_texts)
    prompt = f"""你是一个文档处理专家。以下是同一类文档的 {len(sample_texts)} 份样本内容：

{combined}

请为这类文档生成一个 JSON Schema，用于后续结构化提取。要求：
1. 字段名用英文 snake_case，但每个字段加中文注释（description）。
2. 标明字段类型（string / number / date / enum 等）。
3. 标明是否必填（required）。
4. 如果能推断出枚举值，列在 enum 中。
5. 只输出 JSON，不要额外解释。

输出格式示例：
{
  "schema_name": "invoice",
  "fields": [
    {"name": "invoice_number", "type": "string", "required": true, "description": "发票号码"},
    {"name": "total_amount", "type": "number", "required": true, "description": "总金额"}
  ]
}"""

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.1  # 低温度保证输出稳定
    )
    return json.loads(response.choices[0].message.content)

# 假设上一步聚类结果已存为 dict: cluster_id -> [pdf_filenames]
cluster_map = {
    0: ["invoice_001.pdf", "invoice_002.pdf", "invoice_003.pdf"],
    1: ["contract_001.pdf", "contract_002.pdf"],
}

doc_dir = "./raw_docs"
schemas = {}

for cluster_id, members in cluster_map.items():
    # 每簇采样 2 份
    samples = members[:2]
    texts = [extract_text_from_pdf(os.path.join(doc_dir, f)) for f in samples]
    schema = generate_schema(texts, str(cluster_id))
    schemas[cluster_id] = schema
    print(f"簇 {cluster_id} Schema:")
    print(json.dumps(schema, indent=2, ensure_ascii=False))