使用BGE M3和openGauss DataVec进行向量生成与存储

BGE M3是一款由BAAI开发的多语言高性能文本嵌入模型，能够将文本转化为语义丰富的高维向量表示。本文将围绕BGE M3和向量数据库openGauss DataVec，介绍如何实现文本向量生成与高效存储。通过这两个工具的结合，能够构建更加智能化的数据检索和处理系统。

注：openGauss DataVec容器化部署详见链接。

案例一： FlagEmbedding + openGauss DataVec

环境准备

  • 安装依赖包

FlagEmbedding是一个专注于检索增强型大语言模型的工具包，提供多种文本嵌入模型和重排序模型，使用bge-m3模型前需要先安装该包，详细教程可以参考huggingface官网。

1. pip3 install -U FlagEmbedding
2. pip3 install psycopg2
  • 加载bge-m3模型

实际使用中，FlagEmbedding框架支持通过指定模型名称从Hugging Face模型库自动加载模型。以下将补充说明离线模型的下载步骤，若选择自动加载模型，可直接跳过此部分内容。

1. git lfs install ; git clone https://www.modelscope.cn/BAAI/bge-m3.git

注意这里需要安装好Git LFS工具后才能下载完整的模型数据，附上git-lfs官网下载地址。

实践

在下面的示例中，我们使用FlagEmbedding中的bge-m3嵌入模型生成向量数据，并将其存储在openGauss DataVec向量数据库中。

1. import psycopg2
2. from psycopg2 import sql
3. from typing import List
4. import numpy as np
5. from FlagEmbedding import BGEM3FlagModel

7. def embedding(text):
8. model = BGEM3FlagModel(model_name_or_path = "BAAI/bge-m3")
9. sentence_vector_dict = model.encode(
10. text,
11. return_dense = True,  # 设置返回dense embedding，默认打开
12. return_sparse = False, # 设置返回sparse embedding，默认关闭
13. return_colbert_vecs = False # 设置返回multi-vector(ColBERT)，默认关闭
14. )
15. return sentence_vector_dict.get("dense_vecs")

17. def create_connection(dbname:str, user:str, password:str, host:str, port:int):
18. conn = psycopg2.connect(
19. dbname = dbname,
20. user = user,
21. password = password,
22. host = host,
23. port = port
24. )
25. cursor = conn.cursor()
26. return conn, cursor

28. def create_table(conn, cursor, table_name:str, dim:int):
29. cursor.execute(
30. sql.SQL(
31. "CREATE TABLE IF NOT EXISTS public.{table_name} (id BIGINT PRIMARY KEY, embedding vector({dim}));"
32. ).format(table_name = sql.Identifier(table_name), dim = sql.Literal(dim))
33. )
34. conn.commit()

36. def insert(conn, cursor, table_name:str, embeddings:List[List[float]], ids:List[int]):
37. data = list(zip(ids, embeddings))
38. cursor.executemany(
39. sql.SQL("INSERT INTO public.{table_name} (id, embedding) VALUES(%s, %s);")
40. .format(table_name = sql.Identifier(table_name)), data
41. )
42. conn.commit()
43. print("数据插入成功！")

45. if __name__ == '__main__':
46. text = "openGauss 是一款开源数据库"
47. emb = embedding(text)
48. dimensions = len(emb)
49. print("text : {}, embedding dim : {}, embedding : {} ...".format(text, dimensions, emb[:10]))

51. conn, cursor = create_connection("testdb", "test_user", YourPassword, "localhost", 5432)
52. create_table(conn, cursor, "test_table1", dimensions)
53. insert(conn, cursor, "test_table1", [emb.tolist()], [0])

输出结果如下：

1. text : openGauss 是一款开源数据库, embedding dim : 768, enbedding : [-0.05427849 -0.02701874 -0.05441538 0.0294214 -0.01936925 -0.00815862 0.01310737 -0.0480913 0.01261776 0.2954952] ...
2. 数据插入成功！

openGauss DataVec的具体使用可以参考Python SDK对接向量数据库

案例二：ollama + openGauss DataVec

环境准备

  • 加载模型

ollama安装可以参考openGauss-RAG实践

1. ollama pull bge-m3
  • 检验
1. ollama list

3. NAME              ID             SIZE    MODIFIED
4. bge-m3:latest     790764642607   1.2GB   18 minutes ago

实践

在下面的示例中，我们使用ollama中的bge-m3嵌入模型生成向量数据，并将其存储在openGauss DataVec向量数据库中。

1. import ollama
2. import psycopg2
3. from psycopg2 import sql
4. from typing import List

6. def embedding(text):
7. vector = ollama.embeddings(model="bge-m3", prompt=text)
8. return vector["embedding"]

10. def create_connection(dbname:str, user:str, password:str, host:str, port:int):
11. conn = psycopg2.connect(
12. dbname = dbname,
13. user = user,
14. password = password,
15. host = host,
16. port = port
17. )
18. cursor = conn.cursor()
19. return conn, cursor

21. def create_table(conn, cursor, table_name:str, dim:int):
22. cursor.execute(
23. sql.SQL(
24. "CREATE TABLE IF NOT EXISTS public.{table_name} (id BIGINT PRIMARY KEY, embedding vector({dim}));"
25. ).format(table_name = sql.Identifier(table_name), dim = sql.Literal(dim))
26. )
27. conn.commit()

29. def insert(conn, cursor, table_name:str, embeddings:List[List[float]], ids:List[int]):
30. data = list(zip(ids, embeddings))
31. cursor.executemany(
32. sql.SQL("INSERT INTO public.{table_name} (id, embedding) VALUES(%s, %s);")
33. .format(table_name = sql.Identifier(table_name)), data
34. )
35. conn.commit()
36. print("数据插入成功！")

38. if __name__ == '__main__':
39. text = "openGauss 是一款开源数据库"
40. emb = embedding(text)
41. dimensions = len(emb)
42. print("text : {}, embedding dim : {}, embedding : {} ...".format(text, dimensions, emb[:10]))

44. conn, cursor = create_connection("testdb", "test_user", YourPassword, "localhost", 5432)
45. create_table(conn, cursor, "test_table1", dimensions)
46. insert(conn, cursor, "test_table1", [emb], [0])

输出结果如下：

2. text : openGauss 是一款开源数据库, embedding dim : 768, enbedding : [-0.5359194278717041, 1.3424185514450073, -3.524909734725952, -1.0017194747924805, -0.1950572431087494, 0.28160029649734497, -0.473337858915329, 0.08056074380874634, -0.22012852132320404, -0.9982725977897644] ...
3. 数据插入成功！

openGauss DataVec的具体使用可以参考Python SDK对接向量数据库