使用nomic-embed-text和openGauss DataVec进行向量化搜索

nomic-embed-text是一个专门用于文本转化为高维向量表示的高性能嵌入模型,本文将介绍如何通过nomic-embed-text和openGauss DataVec轻松实现从文本到向量的转化,并基于语义的相似性快速进行搜索操作。

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

环境准备

  • 加载模型

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

1. ollama pull nomic-embed-text
  • 检验
1. ollama list

3. NAME                        ID             SIZE    MODIFIED
4. nomic-embed-text:latest     0a109f422b47   274MB   18 minutes ago

实践

在下面的示例中,我们使用ollama中的nomic-embed-text嵌入模型生成向量数据,并将其存储在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="nomic-embed-text", 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, content text, 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], contents:List[str]):
30. data = list(zip(ids, embeddings, contents))
31. cursor.executemany(
32. sql.SQL("INSERT INTO public.{table_name} (id, embedding, content) VALUES(%s, %s, %s);")
33. .format(table_name = sql.Identifier(table_name)), data
34. )
35. conn.commit()
36. print("数据插入成功!")

39. texts = ["openGauss 是一款开源数据库", "DataVec是一个基于openGauss的向量数据库"]
40. embs = [embedding(t) for t in texts] # 生成底库向量数据
41. dimensions = len(embs[0])
42. ids = [i for i in range(len(embs))]
43. print("text : {}, embedding dim : {}, embedding : {} ...".format(text[0], dimensions, embs[:10]))

45. # 插入数据
46. conn, cursor = create_connection("testdb", "test_user", xxxxxx, "localhost", 5432)
47. create_table(conn, cursor, "test_table1", dimensions)
48. insert(conn, cursor, "test_table1", embs, ids, texts)

输出结果如下:

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

准备好数据后,我们就可以输入查询文本在向量数据库中进行近似查询。

1. def select(conn, cursor, table_name:str, queries:List[List[float]], topk:int):
2. ids = []
3. contents = []
4. for emb in queries:
5. cursor.execute(
6. sql.SQL(
7. "SELECT * FROM public.{table_name} ORDER BY embedding <-> %s::vector LIMIT %s::int;"
8. ).format(table_name = sql.Identifier(table_name)), (emb, topk)
9. )
10. conn.commit()
11. result = cursor.fetchall()
12. ids.append([int(i[0]) for i in result])
13. contents.append([i[1] for i in result])
14. return ids, contents

16. # 生成请求向量数据
17. query = "openGauss数据库是什么"
18. q_emb = embedding(query)

20. # 近似查询
21. ids, contents = select(conn, cursor, "test_table1", [q_emb], 1)
22. print(f"id : {ids[0]}, contents: {contents[0]}")

输出结果如下:

1. id: [0], contents: ['openGauss 是一款开源数据库']

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