# pip install -U psycopgimportpsycopgURL="postgresql://postgres:mysecretpassword@localhost:5432/postgres"withpsycopg.connect(URL)asconn:conn.execute("DROP TABLE IF EXISTS documents;")conn.execute(f"""
CREATE TABLE documents (
id SERIAL PRIMARY KEY,
text TEXT NOT NULL,
dense vector(1024) NOT NULL,
sparse svector({len(model.tokenizer)}) NOT NULL);
""")
表创建完成后,我们可以将嵌入向量插入表中,并为两个向量列分别创建索引:
# pip install -U pgvecto_rsimportpsycopgfrompgvecto_rs.psycopgimportregister_vectorURL="postgresql://postgres:mysecretpassword@localhost:5432/postgres"withpsycopg.connect(URL)asconn:register_vector(conn)fortext,dense,sparse_ind,sparse_valinzip(sentences,dense_embedding,sparse_indices,sparse_values):conn.execute("INSERT INTO documents (text, dense, sparse) VALUES (%s, %s, to_svector(%s, %s::real[]::int[], %s::real[]));",(text,dense,len(model.tokenizer),sparse_ind,sparse_val),)conn.execute("""
CREATE INDEX ON documents
USING vectors (sparse svector_dot_ops)
WITH (options = \"[indexing.hnsw]\");
""")conn.execute("""
CREATE INDEX ON documents
USING vectors (dense vector_l2_ops)
WITH (options = \"[indexing.hnsw]\");
""")
dense_embedding,sparse_indices,sparse_values=embedding("The text you want to search...")withpsycopg.connect(URL)asconn:register_vector(conn)cur=conn.execute("SELECT text FROM documents ORDER BY sparse <#> to_svector(%s, %s::real[]::int[], %s::real[]) LIMIT 50;",(len(model.tokenizer),sparse_indices,sparse_values),)sparse_result=cur.fetchall()cur=conn.execute("SELECT text FROM documents ORDER BY dense <-> %s LIMIT 50;",(dense_embedding,),)dense_result=cur.fetchall()# 合并稀疏向量与稠密向量的搜索结果mix_text=set([r[0]forrindense_result]).union([r[0]forrinsparse_result])
fromFlagEmbeddingimportFlagRerankerreranker=FlagReranker('BAAI/bge-reranker-large',use_fp16=True)# 获取相关性得分(得分越高表示相关性越强)scores=reranker.compute_score([["The text you want to search...",candidate]forcandidateinmix_text])# 对结合了稠密嵌入与稀疏嵌入检索出的候选文本进行重排序reranked_text=[tfor_,tinsorted(zip(scores,mix_text))]