1. 引言

在当今信息爆炸的时代，如何高效地从海量文本数据中检索出相关信息，是许多应用面临的核心挑战。传统的基于关键词的检索方式已难以满足语义搜索的需求。本文将介绍如何结合 TextGen（文本生成工具）、all-mpnet-base-v2 嵌入模型以及 PostgreSQL 数据库的 PGVector 扩展，构建一个完整的向量检索系统。该系统能够将文本数据转换为向量并存储，进而实现基于语义相似度的高效检索。

2. 核心组件介绍

2.1 TextGen

TextGen 是一个用于文本生成和处理的工具或框架。在本方案中，它的主要角色是处理和准备待索引的原始文本数据，例如进行文本清洗、分块或格式化，为后续的向量化步骤做好准备。

(TextGen 安装教程：TextGen for Local LLM-CSDN博客)

2.2 all-mpnet-base-v2 模型

all-mpnet-base-v2 是 Sentence Transformers 库中一个高性能的句子嵌入模型。它基于 MPNet 架构，在多种语义文本相似度（STS）任务上表现出色，能够将句子或段落映射为 768 维的稠密向量。这些向量能够很好地捕捉文本的语义信息，使得语义相似的文本在向量空间中也彼此接近。

(模型安装教程：TextGen 启用 embedding 并安装 all-mpnet-base-v2 模型教程-CSDN博客)

2.3 PostgreSQL 与 PGVector

PostgreSQL 是一款强大的开源关系型数据库。PGVector 是其一个开源扩展，用于存储和检索向量数据。它支持多种向量相似度计算方式（如 L2 距离、内积、余弦相似度），并提供了高效的索引（如 IVFFlat, HNSW）来加速大规模向量检索。

3. 系统架构与工作流程

整个系统的工作流程可以分为以下几个步骤：

  1. 数据准备：使用 TextGen 处理原始文本，生成适合嵌入模型输入的文本块。
  2. 向量化：使用 all-mpnet-base-v2 模型将文本块转换为向量。
  3. 存储：将文本内容及其对应的向量存储到启用了 PGVector 扩展的 PostgreSQL 数据库中。
  4. 检索：将用户的查询语句同样转换为向量，并在数据库中进行相似度搜索，返回最相关的文本结果。

4. 实践步骤

4.1 环境准备

首先，确保你的环境已安装必要的 Python 库和 PostgreSQL 扩展。

# 安装 Python 依赖
pip install sentence-transformers psycopg2-binary pgvector
在 PostgreSQL 中创建 PGVector 扩展 (需要在数据库内执行)
CREATE EXTENSION IF NOT EXISTS vector;

4.2 数据库表设计

创建一个用于存储文本和向量的表。

在 user_data/extensions/rag/index_to_pg.py 中由程序自动创建。

4.3 文本处理与向量化

使用 TextGen 和 all-mpnet-base-v2 模型生成向量。

user_data/extensions/rag/index_to_pg.py

# user_data/extensions/rag/index_to_pg.py
"""
索引脚本：支持 TXT/PDF/Excel 的增量 upsert，含自动建表与索引。
"""
import os
import hashlib
from pathlib import Path
from tqdm import tqdm
import psycopg2
from psycopg2.extras import Json
import pdfplumber
import pandas as pd
from modules.api.embeddings import get_embeddings, get_embeddings_model_name
try:
    from pgvector.psycopg2 import register_vector, Vector
except ImportError:
    try:
        from pgvector.psycopg2 import register_vector, vector as Vector
    except Exception:
        try:
            from pgvector import Vector
            from pgvector.psycopg2 import register_vector
        except Exception:
            register_vector = None
            Vector = None

USE_PGVECTOR_WRAPPER = callable(Vector)

PG_CONN = os.environ.get("PG_CONN", "host=active.aiengineerdev.service.paas.wistron.com port=15237 dbname=dlx user=dxlab password=wistron888")
DOC_DIR = Path(os.environ.get("PG_DOC_DIR", "knowledge"))
BATCH = int(os.environ.get("PG_BATCH", 64))
CHUNK_WORDS = int(os.environ.get("CHUNK_WORDS", 400))
CHUNK_OVERLAP = int(os.environ.get("CHUNK_OVERLAP", 50))

def chunks_of_text(text, chunk_words=CHUNK_WORDS, overlap=CHUNK_OVERLAP):
    words = text.split()
    i = 0
    while i < len(words):
        chunk = " ".join(words[i:i+chunk_words])
        yield chunk
        i += max(1, chunk_words - overlap)

def hash_text(s: str) -> str:
    return hashlib.sha256(s.encode("utf-8")).hexdigest()

def extract_text_from_pdf(path: Path) -> str:
    texts = []
    try:
        with pdfplumber.open(path) as pdf:
            for page in pdf.pages:
                t = page.extract_text()
                if t:
                    texts.append(t)
    except Exception:
        return ""
    return "\n".join(texts)

def extract_texts_from_excel(path: Path):
    parts = []
    try:
        xls = pd.read_excel(path, sheet_name=None, engine="openpyxl")
    except Exception:
        xls = pd.read_excel(path, sheet_name=None)
    for sheet_name, df in xls.items():
        for idx, row in df.iterrows():
            vals = [str(v) for v in row.values if pd.notna(v)]
            if vals:
                parts.append((sheet_name, int(idx), " ".join(vals)))
    return parts

def detect_embedding_dim():
    sample = ["hello world"]
    embs = get_embeddings(sample)
    first = embs[0]
    try:
        return len(first)
    except Exception:
        # fallback to env or 1536
        return int(os.environ.get("PG_VECTOR_DIM", 1536))

def create_table_if_needed(cur, dim):
    # create extension and table if not exists
    cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
    cur.execute(f"""
    CREATE TABLE IF NOT EXISTS documents (
        id SERIAL PRIMARY KEY,
        doc_id TEXT NOT NULL,
        chunk_index INTEGER,
        text TEXT,
        metadata JSONB,
        embedding vector({dim}),
        chunk_hash TEXT
    );
    """)
    cur.execute("CREATE UNIQUE INDEX IF NOT EXISTS documents_docid_chunkhash_idx ON documents (doc_id, chunk_hash);")
    # optional ivfflat index (uncomment if you want and tune lists)
    # cur.execute(f"CREATE INDEX IF NOT EXISTS documents_embedding_idx ON documents USING ivfflat (embedding) WITH (lists = 100);")

def upsert_chunks_for_doc(conn, cur, doc_id, items):
    # items: list of dict {chunk_index, text, metadata, chunk_hash}
    cur.execute("SELECT chunk_hash FROM documents WHERE doc_id=%s", (doc_id,))
    existing = {r[0] for r in cur.fetchall() if r[0] is not None}

    new_items = [it for it in items if it["chunk_hash"] not in existing]
    kept_hashes = {it["chunk_hash"] for it in items}

    to_delete = list(existing - kept_hashes)
    if to_delete:
        cur.execute("DELETE FROM documents WHERE doc_id=%s AND chunk_hash = ANY(%s)", (doc_id, to_delete))
        conn.commit()

    # 根据环境选择 SQL 占位与 embedding 值的构造方式
    if USE_PGVECTOR_WRAPPER:
        insert_sql = """
            INSERT INTO documents (doc_id, chunk_index, text, metadata, embedding, chunk_hash)
            VALUES (%s,%s,%s,%s,%s,%s)
            ON CONFLICT (doc_id, chunk_hash) DO UPDATE
              SET text = EXCLUDED.text,
                  metadata = EXCLUDED.metadata,
                  embedding = EXCLUDED.embedding,
                  chunk_index = EXCLUDED.chunk_index
        """
    else:
        # 将 embedding 作为向量字面量字符串，并在 SQL 中 cast 为 vector
        insert_sql = """
            INSERT INTO documents (doc_id, chunk_index, text, metadata, embedding, chunk_hash)
            VALUES (%s,%s,%s,%s,%s::vector,%s)
            ON CONFLICT (doc_id, chunk_hash) DO UPDATE
              SET text = EXCLUDED.text,
                  metadata = EXCLUDED.metadata,
                  embedding = EXCLUDED.embedding,
                  chunk_index = EXCLUDED.chunk_index
        """

    for i in range(0, len(new_items), BATCH):
        batch = new_items[i:i+BATCH]
        texts = [it["text"] for it in batch]
        embs = get_embeddings(texts)
        records = []
        for j, emb in enumerate(embs):
            it = batch[j]
            if USE_PGVECTOR_WRAPPER:
                emb_val = Vector(list(map(float, emb)))
                records.append((
                    doc_id,
                    it["chunk_index"],
                    it["text"],
                    Json(it["metadata"]),
                    emb_val,
                    it["chunk_hash"],
                ))
            else:
                # 构造向量字面量，例如: "[0.1,0.2,...]"
                emb_list = list(map(float, emb))
                emb_str = "[" + ",".join(map(lambda x: repr(float(x)), emb_list)) + "]"
                records.append((
                    doc_id,
                    it["chunk_index"],
                    it["text"],
                    Json(it["metadata"]),
                    emb_str,
                    it["chunk_hash"],
                ))
        cur.executemany(insert_sql, records)
        conn.commit()

def process_file(conn, cur, path: Path):
    doc_id = str(path.relative_to(DOC_DIR))
    ext = path.suffix.lower()
    items = []

    if ext == ".txt":
        txt = path.read_text(encoding="utf-8", errors="ignore")
        for idx, chunk in enumerate(chunks_of_text(txt)):
            items.append({
                "chunk_index": idx,
                "text": chunk,
                "metadata": {"source": doc_id, "type": "text"},
                "chunk_hash": hash_text(chunk),
            })
    elif ext == ".pdf":
        txt = extract_text_from_pdf(path)
        for idx, chunk in enumerate(chunks_of_text(txt)):
            items.append({
                "chunk_index": idx,
                "text": chunk,
                "metadata": {"source": doc_id, "type": "pdf"},
                "chunk_hash": hash_text(chunk),
            })
    elif ext in (".xls", ".xlsx"):
        parts = extract_texts_from_excel(path)
        seq = 0
        for sheet, row_idx, text in parts:
            for chunk in chunks_of_text(text):
                items.append({
                    "chunk_index": seq,
                    "text": chunk,
                    "metadata": {"source": doc_id, "type": "excel", "sheet": sheet, "row": row_idx},
                    "chunk_hash": hash_text(chunk),
                })
                seq += 1
    else:
        return

    if items:
        upsert_chunks_for_doc(conn, cur, doc_id, items)

def index_all():
    conn = psycopg2.connect(PG_CONN)
    register_vector(conn)
    cur = conn.cursor()

    dim = int(os.environ.get("PG_VECTOR_DIM") or detect_embedding_dim())
    create_table_if_needed(cur, dim)
    conn.commit()

    files = list(DOC_DIR.glob("**/*"))
    files = [f for f in files if f.suffix.lower() in (".txt", ".pdf", ".xls", ".xlsx")]
    print("Embedding model:", get_embeddings_model_name())
    for f in tqdm(files, desc="files"):
        try:
            process_file(conn, cur, f)
        except Exception as e:
            print("Error processing", f, e)

    cur.close()
    conn.close()
    print("Indexing complete.")

if __name__ == "__main__":
    index_all()

user_data/extensions/rag/1.index_to_pg.sh

export root_dir=/home/ai/textgen-4.9
export OPENEDAI_EMBEDDING_MODEL="/home/ai/textgen-4.9/user_data/models/embedding/all-mpnet-base-v2"
export OPENEDAI_EMBEDDING_DEVICE="auto"   # "cpu" or "cuda" or auto

export PYTHONPATH=$root_dir/app
$root_dir/app/portable_env/bin/python $root_dir/user_data/extensions/rag/index_to_pg.py


#檢查CUDA是否可用
# $root_dir/app/portable_env/bin/python -c "import torch; print('cuda available:', torch.cuda.is_available()); print(torch.__version__)"

执行 ./1.index_to_pg.sh：

 Loaded embedding model: /home/ai/textgen-4.9/user_data/models/embedding/all-mpnet-base-v2
Embedding model: /home/ai/textgen-4.9/user_data/models/embedding/all-mpnet-base-v2
files: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00,  3.00it/s]
Indexing complete.
 

4.4 向量检索

当用户输入查询时，将其转换为向量并检索最相似的文档。

user_data/extensions/rag/retriever_pg.py

# user_data/extensions/rag/retriever_pg.py
import os
import psycopg2
import json
from pgvector.psycopg2 import register_vector
try:
    from pgvector import Vector
except Exception:
    try:
        from pgvector.psycopg2 import vector as Vector
    except Exception:
        Vector = None

from modules.api.embeddings import get_embeddings
from modules.logging_colors import logger

PG_CONN = os.environ.get("PG_CONN", "host=active.aiengineerdev.service.paas.wistron.com port=15237 dbname=dlx user=dxlab password=wistron888")

def _connect():
    conn = psycopg2.connect(PG_CONN)
    try:
        register_vector(conn)
    except Exception:
        pass
    return conn

def retrieve(query, top_k=5):
    conn = _connect()
    cur = conn.cursor()
    vec = get_embeddings([query])[0]  # numpy array
    # 如果有 pgvector 的 Vector 包装器则使用它，否则传原生浮点列表
    if Vector:
        v = Vector(list(map(float, vec)))
    else:
        v = list(map(float, vec))

    cur.execute("""
      SELECT id, doc_id, chunk_index, text, metadata, embedding <-> %s AS distance
      FROM documents
      ORDER BY embedding <-> %s
      LIMIT %s
    """, (v, v, top_k))
    rows = cur.fetchall()
    results = []
    for r in rows:
        results.append({
            "id": r[0],
            "doc_id": r[1],
            "chunk_index": r[2],
            "text": r[3],
            "metadata": r[4],
            "score": float(r[5])
        })
    cur.close()
    conn.close()
    try:
        q_preview = str(query).replace("\n", " ")[:100]
        logger.info(f"[rag] retrieve: query='{q_preview}' top_k={top_k} returned={len(results)}")
    except Exception:
        pass
    return results

############ For testing ##########################
#if __name__ == "__main__":
#    results = retrieve("基本(最低)工資實施為何未發布日薪?", top_k=5)
#    print(json.dumps(results, ensure_ascii=False, indent=2))
##################################################

user_data/extensions/rag/2.retriever_pg.sh

# 執行之前需要移除 retriever_pg.py 中的 if __name__ == "__main__":... 註釋

export root_dir=/home/ai/textgen-4.9
export OPENEDAI_EMBEDDING_MODEL="/home/ai/textgen-4.9/user_data/models/embedding/all-mpnet-base-v2"
export OPENEDAI_EMBEDDING_DEVICE="auto"   # "cpu" or "cuda" or auto

export PYTHONPATH=$root_dir/app
$root_dir/app/portable_env/bin/python $root_dir/user_data/extensions/rag/retriever_pg.py

执行 ./2.retriever_pg.sh

Loading weights: 100%|████████████████████████████████████████████████████████████████████████████████████████████| 199/199 [00:00<00:00, 10877.40it/s]
23:51:41-890216 INFO     Loaded embedding model: /home/ai/textgen-4.9/user_data/models/embedding/all-mpnet-base-v2
23:51:42-015077 INFO     [rag] retrieve: query='基本(最低)工資實施為何未發布日薪?' top_k=5 returned=5
[
  {
    "id": 697,
    "doc_id": "labor_qa_dataset.xlsx",
    "chunk_index": 335,
     ...

5. 优化与注意事项

  • 索引选择：PGVector 支持 IVFFlat 和 HNSW 索引。对于大规模数据集，HNSW 通常提供更好的查询性能。
  • 分块策略：TextGen 的分块逻辑至关重要。过小的块可能丢失上下文，过大的块可能降低检索精度。需要根据实际数据调整。
  • 模型选择：all-mpnet-base-v2 是通用性很强的模型。对于特定领域（如医学、法律），可以考虑使用在该领域微调过的嵌入模型。
  • 混合搜索：可以结合向量相似度（语义）和传统关键词匹配（BM25）进行混合检索，以兼顾召回率与精确度。

6. 总结

通过整合 TextGen、all-mpnet-base-v2 和 PostgreSQL (PGVector)，我们构建了一个强大且灵活的向量检索系统。该方案利用了 Sentence Transformers 强大的语义编码能力，以及 PostgreSQL 作为成熟数据库的可靠性与扩展性，能够有效地解决语义搜索、推荐系统、问答机器人等多种场景下的信息检索需求。开发者可以根据具体业务需求，调整文本处理流程、模型和检索参数，以获得最佳效果。

# postgresql # 人工智能 # 自然语言处理 # 目标检测
Logo
https://edu.csdn.net/learn/39067/627173?utm_source=2019755004

汇聚全球AI编程工具，助力开发者即刻编程。

更多推荐

【无标题】

一开始觉得这样很偷懒，后来发现其实比手写还累。你得把需求描述得足够精确AI才能给对的东西，这个过程其实在逼你提升技术表达能力和代码审查能力。还有就是经常发现AI写的代码有隐藏bug，不仔细看根本看不出来。我会先在脑子里把功能逻辑想清楚，这步不能跳。然后用大白话描述给AI，尽量把边界条件和异常处理也说清楚。它生成之后我一行行读，看到不对的就标出来让它改，改完我再读，来回到我满意为止。大家用AI写代码

avatar AI编程社区
cover

Claude Code让普通人也能编程

avatar AI编程社区

Cursor + GitOps：自动化运维新姿势

Cursor与GitOps的结合，并非要取代工程师，而是将工程师从重复、繁琐的配置编写中解放出来，更专注于架构设计与核心业务逻辑。它代表了一种新的协作范式：人类负责定义“做什么”和“为什么”，AI助手高效地完成“怎么做”。掌握这一新姿势，是迈向下一代自动化运维的关键一步。

avatar AI编程社区