Transformer 学习之路 - 向量召回与排序实战
Transformer 学习之路 - 向量召回与排序实战
在自然语言处理（NLP）领域，Transformer 模型已经成为一种革命性的技术。它不仅广泛应用于机器翻译、文本生成等任务，还在信息检索、问答系统等领域展现了强大的能力。本文将结合代码示例，深入探讨 Transformer 在向量召回与排序中的应用，帮助你理解其技术原理并实际应用。
1. 背景与问题
在问答系统中，如何快速准确地找到与用户问题最相关的答案是一个核心挑战。传统的基于关键词匹配的方法往往无法捕捉语义信息，导致召回效果不佳。而 Transformer 模型通过将文本转化为高维向量，能够更好地表达语义信息，从而实现更精准的召回与排序。
2. 技术原理
2.1 向量召回
向量召回的核心思想是将文本（如问题与答案）映射到高维向量空间，然后通过计算向量之间的相似度来找到最相关的结果。具体步骤如下：
- 文本向量化：使用预训练的 Transformer 模型（如 Sentence-BERT）将文本转化为固定长度的向量。
- 向量存储：将生成的向量存储到向量数据库（如 PostgreSQL 的
vector扩展）。 - 相似度计算：通过计算查询向量与存储向量的余弦相似度，找到最相关的结果。
2.2 向量排序
在召回的结果中，可能存在多个相关但质量不同的答案。为了进一步提升结果质量，可以使用排序模型对召回结果进行二次排序。排序模型通过计算查询与候选答案的匹配分数，选择最合适的答案。
3. 代码实现
3.1 读取数据集
首先，我们加载一个法律问答数据集（law_faq.csv），其中包含问题（title）和答案（reply）。
import pandas as pd
data = pd.read_csv('law_faq.csv', encoding='utf-8')
documents = data['title'].tolist()
print(documents[0:5]) # 打印前5个问题
3.2 加载向量模型
使用 sentence-transformers 库加载预训练的向量模型，并将文本转化为向量。
from sentence_transformers import SentenceTransformer
model = SentenceTransformer("TencentBAC/Conan-embedding-v1", cache_folder="cache")
document_vectors = model.encode(documents[0:5], convert_to_numpy=True)
print(document_vectors.shape) # 打印向量维度
3.3 连接 PostgreSQL 数据库
将生成的向量存储到 PostgreSQL 数据库中，并创建向量索引以加速查询。
import psycopg2
connection = psycopg2.connect(
host="127.0.0.1",
port="5432",
user="root",
dbname="test",
password="Ym135168.",
)
connection.autocommit = True
curr = connection.cursor()
# 创建向量扩展和表
curr.execute("CREATE EXTENSION IF NOT EXISTS vector;")
curr.execute("""
CREATE TABLE IF NOT EXISTS public.law_faq (
id SERIAL PRIMARY KEY,
title VARCHAR NOT NULL,
reply VARCHAR NOT NULL,
embedding vector(1792) NOT NULL
);
""")
curr.execute("CREATE INDEX ON public.law_faq USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);")
3.4 向量化数据并写入数据库
将数据集中的文本批量转化为向量，并存储到数据库中。
from tqdm import tqdm
from psycopg2.extras import execute_batch
batch_size = 5
for i in tqdm(range(0, len(documents), batch_size), desc="Encoding Progress"):
batch = documents[i:i + batch_size]
batch_vectors = model.encode(batch, convert_to_numpy=True)
insert_data = [(data.values[i + j][0], data.values[i + j][1], batch_vectors[j].tolist()) for j in range(len(batch))]
execute_batch(curr, "INSERT INTO public.law_faq (title, reply, embedding) VALUES (%s, %s, %s)", insert_data)
3.5 根据向量查询
将用户查询转化为向量，并从数据库中检索最相关的结果。
query = "怎么才算寻衅滋事呢？"
query_vector = model.encode(query, convert_to_numpy=True)
def retrieve_similar(curr, query_vector, top_k=5):
query_vector = query_vector.flatten().tolist()
query_vector_sql = "ARRAY[" + ",".join(map(str, query_vector)) + "]::vector"
curr.execute(f"""
SELECT id, title, reply, embedding
FROM public.law_faq
ORDER BY embedding <-> {query_vector_sql}
LIMIT %s;
""", (top_k,))
return curr.fetchall()
similar_results = retrieve_similar(curr, query_vector)
for result in similar_results:
print(f"问题: {result[1]}\n回复: {result[2]}")
3.6 加载排序模型并打分
使用 FlagEmbedding 库对召回结果进行排序，选择最佳答案。
from FlagEmbedding import FlagReranker
reranker = FlagReranker('BAAI/bge-reranker-base', use_fp16=True, cache_dir="cache")
ranker_data = [[query, i[2]] for i in similar_results]
ranker_results = reranker.compute_score(ranker_data, normalize=True)
max_index = ranker_results.index(max(ranker_results))
print("问题:", query)
print("最佳回复标题:", similar_results[max_index][1])
print("最佳回复内容:", similar_results[max_index][2])
4. 总结
本文通过一个法律问答系统的示例，详细介绍了如何使用 Transformer 模型实现向量召回与排序。通过将文本转化为向量并利用向量数据库进行高效检索，我们能够快速找到与用户问题最相关的答案。此外，通过引入排序模型，进一步提升了结果的质量。希望本文能帮助你深入理解 Transformer 技术的应用，并在实际项目中灵活运用。
如果你对代码实现有任何疑问，欢迎在评论区留言讨论！