AI 工作流专栏 14：Embedding 与文档处理实战

# embedding_intuition.py: 先"看见"Embedding 长什么样
import os
from openai import OpenAI

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

texts = [
    "我爱吃火锅",            # 1
    "火锅是我的最爱",         # 2, 语义和 1 几乎一样
    "今天股市大跌",          # 3, 语义和 1 完全无关
]

resp = client.embeddings.create(
    model="text-embedding-3-small",
    input=texts,
)
v1, v2, v3 = [d.embedding for d in resp.data]

# 余弦相似度（手写版，理解原理）
import math
def cosine(a, b):
    dot = sum(x*y for x, y in zip(a, b))
    na = math.sqrt(sum(x*x for x in a))
    nb = math.sqrt(sum(y*y for y in b))
    return dot / (na * nb)

print(f"火锅 vs 火锅: {cosine(v1, v2):.3f}")   # ≈ 0.92，很相似
print(f"火锅 vs 股市: {cosine(v1, v3):.3f}")   # ≈ 0.45，不相关

# embedding_openai.py: 用 OpenAI 跑 Embedding
import os
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

def embed_openai(texts: list[str], model="text-embedding-3-small") -> list[list[float]]:
    """批量 Embedding，OpenAI 单次最多 2048 条"""
    all_vecs = []
    for i in range(0, len(texts), 2048):
        batch = texts[i:i+2048]
        resp = client.embeddings.create(model=model, input=batch)
        # 返回顺序和输入一致，保险起见按 index 排个序
        all_vecs.extend([d.embedding for d in sorted(resp.data, key=lambda x: x.index)])
    return all_vecs

if __name__ == "__main__":
    vecs = embed_openai(["RAG 是检索增强生成", "Retrieval-Augmented Generation"])
    print(f"维度: {len(vecs[0])}, 第一条前 5 维: {vecs[0][:5]}")

# embedding_bge.py: 用 sentence-transformers 本地跑 BGE-M3
# 先装：pip install sentence-transformers
from sentence_transformers import SentenceTransformer
import numpy as np

# 首次运行会自动下载模型（约 2GB），之后从本地加载
model = SentenceTransformer(
    "BAAI/bge-m3",                # 智源 BGE-M3
    device="cpu",                 # 没 GPU 就用 cpu；有 cuda 写 "cuda"
)

def embed_bge(texts: list[str]) -> np.ndarray:
    """本地批量 Embedding，返回 numpy 数组"""
    vecs = model.encode(
        texts,
        batch_size=32,
        normalize_embeddings=True,  # 归一化后点积 = 余弦相似度，检索更快
        show_progress_bar=False,
    )
    return vecs

if __name__ == "__main__":
    vecs = embed_bge(["RAG 是检索增强生成", "Retrieval-Augmented Generation"])
    print(f"维度: {vecs.shape[1]}")  # 1024
    print(f"两句话余弦相似度: {float(vecs[0] @ vecs[1]):.3f}")  # 归一化后点积即余弦

# parse_pdf.py: 用 pdfplumber 提取 PDF
# 先装：pip install pdfplumber
import pdfplumber

def extract_pdf(pdf_path: str):
    """提取 PDF 全部文字 + 表格，返回 (full_text, tables)"""
    full_text = []
    all_tables = []
    with pdfplumber.open(pdf_path) as pdf:
        for page_idx, page in enumerate(pdf.pages):
            # 1) 提取纯文字
            text = page.extract_text() or ""
            full_text.append(f"\n--- 第 {page_idx+1} 页 ---\n{text}")

            # 2) 提取表格（pdfplumber 的杀手锏）
            for table in page.extract_tables():
                all_tables.append({"page": page_idx + 1, "rows": table})

    return "\n".join(full_text), all_tables

if __name__ == "__main__":
    text, tables = extract_pdf("sample.pdf")
    print(f"文字总长度: {len(text)} 字符")
    print(f"表格数量: {len(tables)}")
    if tables:
        print(f"第一个表格第一行: {tables[0]['rows'][0]}")

# parse_unstructured.py: Unstructured 万能解析
# 先装：pip install "unstructured[all-docs]"
from unstructured.partition.auto import partition

elements = partition("sample.pdf")   # 自动识别格式
for el in elements:
    print(el.category, ":", str(el)[:80])
# 输出形如：
# Title : 第一章 项目背景
# Narrative : 本项目旨在...
# Table : [['月份','销量'], ['1月','1200'], ...]

# parse_docx.py
# pip install python-docx
from docx import Document

def extract_docx(path: str) -> str:
    doc = Document(path)
    parts = []
    for para in doc.paragraphs:
        # 保留标题层级信息（Heading 1/2/3）
        style = para.style.name if para.style else "Normal"
        prefix = ""
        if style.startswith("Heading"):
            level = style.replace("Heading ", "")
            prefix = "#" * int(level) + " " if level.isdigit() else ""
        if para.text.strip():
            parts.append(prefix + para.text.strip())
    return "\n\n".join(parts)

# parse_xlsx.py
# pip install openpyxl
from openpyxl import load_workbook

def extract_xlsx(path: str) -> list[dict]:
    """每个 sheet 转成一个 Markdown 表格字符串"""
    wb = load_workbook(path, data_only=True)  # data_only 取计算后的值
    sheets = []
    for ws in wb.worksheets:
        rows = list(ws.iter_rows(values_only=True))
        if not rows:
            continue
        # 转成 Markdown 表格
        header = rows[0]
        md = ["| " + " | ".join(str(c) if c is not None else "" for c in header) + " |"]
        md.append("| " + " | ".join("---" for _ in header) + " |")
        for row in rows[1:]:
            md.append("| " + " | ".join(str(c) if c is not None else "" for c in row) + " |")
        sheets.append({"sheet": ws.title, "markdown": "\n".join(md)})
    return sheets

# parse_pptx.py
# pip install python-pptx
from pptx import Presentation

def extract_pptx(path: str) -> str:
    prs = Presentation(path)
    parts = []
    for i, slide in enumerate(prs.slides, 1):
        parts.append(f"\n## 幻灯片 {i}")
        for shape in slide.shapes:
            if shape.has_text_frame:
                for para in shape.text_frame.paragraphs:
                    txt = para.text.strip()
                    if txt:
                        parts.append(txt)
            elif shape.has_table:   # PPT 里的表格
                for row in shape.table.rows:
                    parts.append(" | ".join(c.text for c in row.cells))
    return "\n".join(parts)

# parse_html_bs4.py
# pip install beautifulsoup4 lxml
import requests
from bs4 import BeautifulSoup

def extract_html_manual(url: str) -> str:
    html = requests.get(url, timeout=10).text
    soup = BeautifulSoup(html, "lxml")

    # 1) 砍掉噪声标签
    for tag in soup(["script", "style", "nav", "footer", "header", "aside", "form"]):
        tag.decompose()

    # 2) 拿正文（article / main / 最长的 div，经验法则）
    main = soup.find("article") or soup.find("main") or soup.body
    text = main.get_text(separator="\n", strip=True) if main else ""

    # 3) 合并多余空行
    lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
    return "\n".join(lines)

# parse_html_trafilatura.py
# pip install trafilatura
import trafilatura

def extract_html_auto(url: str) -> str:
    html = trafilatura.fetch_url(url)
    # include_tables=True 保留表格，include_links=True 保留链接
    text = trafilatura.extract(
        html,
        include_tables=True,
        include_links=False,
        favor_recall=True,        # 偏向召回，多抓点正文
    )
    return text or ""

| 月份 | 销量 | 同比 |
|------|------|------|
| 1月  | 1200 | +15% |
| 2月  | 1350 | +12% |

[{"月份": "1月", "销量": 1200, "同比": "+15%"},
 {"月份": "2月", "销量": 1350, "同比": "+12%"}]

# ocr_paddle.py
# pip install paddlepaddle paddleocr
from paddleocr import PaddleOCR

ocr = PaddleOCR(use_angle_cls=True, lang="ch")  # 中文
result = ocr.ocr("invoice.png", cls=True)
for line in result[0]:
    print(line[1][0])   # line[1][0] 是识别出的文字

# image_to_text_vlm.py: 视觉大模型读图
import os, base64
from openai import OpenAI

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

def image_to_text(image_path: str) -> str:
    with open(image_path, "rb") as f:
        b64 = base64.b64encode(f.read()).decode()

    resp = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{
            "role": "user",
            "content": [
                {"type": "text", "text": "详细描述这张图的内容，尤其是图表中的数据和趋势，输出纯文本，便于检索。"},
                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}},
            ],
        }],
    )
    return resp.choices[0].message.content

# clean_text.py: 文档清洗通用函数
import re

def clean_text(text: str) -> str:
    # 1) 统一换行符（Windows 的 \r\n → \n）
    text = text.replace("\r\n", "\n").replace("\r", "\n")

    # 2) 去页眉页脚的常见模式（页码、"第 X 页"）
    text = re.sub(r"-?\s*\d+\s*-?\s*\n", "\n", text)       # 单独一行的页码
    text = re.sub(r"第\s*\d+\s*页.*?\n", "", text)          # 中文页码
    text = re.sub(r"Page\s+\d+\s+of\s+\d+", "", text)       # 英文页码

    # 3) 去乱码（非打印字符、Unicode 替换符）
    text = text.replace("\ufffd", "")                       # 替换符
    text = re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f]", "", text)  # 控制字符

    # 4) 合并多余空白：连续空格 → 单空格，连续空行 → 单空行
    text = re.sub(r"[ \t]+", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)

    # 5) 修复断行：行尾是中文或小写字母且下一行不是新段落，拼回去
    text = re.sub(r"([^\n。！？.!?])\n([^\n])", r"\1\2", text)

    return text.strip()

def chunk_fixed(text: str, size=500, overlap=50) -> list[str]:
    chunks = []
    start = 0
    while start < len(text):
        chunks.append(text[start:start+size])
        start += size - overlap      # 滑动窗口，留 overlap 重叠
    return chunks

# pip install langchain-text-splitters
from langchain_text_splitters import MarkdownHeaderTextSplitter

def chunk_by_header(md_text: str) -> list[str]:
    splitter = MarkdownHeaderTextSplitter(
        headers_to_split_on=[("#", "h1"), ("##", "h2"), ("###", "h3")]
    )
    docs = splitter.split_text(md_text)
    # 每段 doc 自带 metadata（属于哪个标题），对检索超有用
    return [f"[{d.metadata}]\n{d.page_content}" for d in docs]

from langchain_text_splitters import RecursiveCharacterTextSplitter

def chunk_recursive(text: str, chunk_size=500, overlap=80) -> list[str]:
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=overlap,
        separators=["\n\n", "\n", "。", ".", " ", ""],  # 按优先级递归
    )
    return splitter.split_text(text)

from langchain_text_splitters import RecursiveCharacterTextSplitter, Language

code_splitter = RecursiveCharacterTextSplitter.from_language(
    language=Language.PYTHON,
    chunk_size=800,
    chunk_overlap=100,
)
code_chunks = code_splitter.split_text(python_source)

# Python 依赖
pip install pdfplumber langchain-text-splitters sentence-transformers \
            psycopg[binary] pgvector python-dotenv

# PostgreSQL 要装 pgvector 扩展
# docker run -d --name pg -e POSTGRES_PASSWORD=postgres -p 5432:5432 pgvector/pgvector:pg16

# rag_pipeline.py: RAG 文档处理完整流水线
"""
把一份 PDF：解析 → 清洗 → 切片 → Embedding → 存进 pgvector
跑通后，检索端就能直接用这个库。
"""
import os, re, hashlib
import pdfplumber
import psycopg
from dotenv import load_dotenv
from sentence_transformers import SentenceTransformer
from langchain_text_splitters import RecursiveCharacterTextSplitter

load_dotenv()

# ---------- 第 1 步：解析 PDF ----------
def parse_pdf(pdf_path: str) -> str:
    parts = []
    with pdfplumber.open(pdf_path) as pdf:
        for i, page in enumerate(pdf.pages, 1):
            text = page.extract_text() or ""
            parts.append(f"\n## 第 {i} 页\n{text}")
            for table in page.extract_tables():
                md = table_to_markdown(table)
                parts.append(md)
    return "\n\n".join(parts)

def table_to_markdown(table) -> str:
    """把 pdfplumber 的二维列表转 Markdown 表格"""
    if not table:
        return ""
    rows = [[(c or "").strip() for c in row] for row in table]
    header = rows[0]
    md = ["| " + " | ".join(header) + " |",
          "| " + " | ".join("---" for _ in header) + " |"]
    for row in rows[1:]:
        md.append("| " + " | ".join(row) + " |")
    return "\n".join(md)

# ---------- 第 2 步：清洗 ----------
def clean_text(text: str) -> str:
    text = text.replace("\r\n", "\n").replace("\r", "\n")
    text = re.sub(r"-?\s*\d+\s*-?\s*\n", "\n", text)       # 页码
    text = text.replace("\ufffd", "")
    text = re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f]", "", text)
    text = re.sub(r"[ \t]+", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()

# ---------- 第 3 步：切片 ----------
def chunk_text(text: str, chunk_size=500, overlap=80) -> list[str]:
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=overlap,
        separators=["\n\n", "\n", "。", ".", " ", ""],
    )
    return splitter.split_text(text)

# ---------- 第 4 步：Embedding（本地 BGE-M3） ----------
_model = None
def get_model():
    global _model
    if _model is None:
        _model = SentenceTransformer("BAAI/bge-m3", device="cpu")
    return _model

def embed(texts: list[str]) -> list[list[float]]:
    model = get_model()
    vecs = model.encode(texts, batch_size=32,
                        normalize_embeddings=True, show_progress_bar=True)
    return vecs.tolist()

# ---------- 第 5 步：存 pgvector ----------
def ensure_table(conn):
    with conn.cursor() as cur:
        cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
        cur.execute("""
            CREATE TABLE IF NOT EXISTS docs (
                id        TEXT PRIMARY KEY,
                source    TEXT,
                page      INT,
                chunk_idx INT,
                content   TEXT,
                embedding vector(1024)
            );
        """)
        conn.commit()

def store_chunks(chunks, vecs, source, conn):
    with conn.cursor() as cur:
        for i, (chunk, vec) in enumerate(zip(chunks, vecs)):
            doc_id = hashlib.md5(f"{source}:{i}:{chunk[:32]}".encode()).hexdigest()
            page = 1  # 简化：实际可从切片时记录的页码取
            vec_str = "[" + ",".join(f"{x:.6f}" for x in vec) + "]"
            cur.execute("""
                INSERT INTO docs (id, source, page, chunk_idx, content, embedding)
                VALUES (%s, %s, %s, %s, %s, %s)
                ON CONFLICT (id) DO NOTHING;
            """, (doc_id, source, page, i, chunk, vec_str))
        conn.commit()

# ---------- 主流程 ----------
def main():
    pdf_path = "sample.pdf"
    source_name = os.path.basename(pdf_path)

    print("[1/5] 解析 PDF...")
    raw = parse_pdf(pdf_path)
    print(f"  原始文本 {len(raw)} 字符")

    print("[2/5] 清洗文本...")
    cleaned = clean_text(raw)
    print(f"  清洗后 {len(cleaned)} 字符")

    print("[3/5] 切片...")
    chunks = chunk_text(cleaned, chunk_size=500, overlap=80)
    print(f"  切出 {len(chunks)} 个 chunk")

    print("[4/5] Embedding (BGE-M3)...")
    vecs = embed(chunks)
    print(f"  向量维度 {len(vecs[0])}")

    print("[5/5] 写入 pgvector...")
    dsn = os.getenv("PG_DSN", "postgresql://postgres:postgres@localhost:5432/postgres")
    with psycopg.connect(dsn) as conn:
        ensure_table(conn)
        store_chunks(chunks, vecs, source_name, conn)
    print(f"完成！{len(chunks)} 个 chunk 已入库，下一步就能检索了。")

if __name__ == "__main__":
    main()