Browse Source

使用ark的向量模型,完成文档的向量化

zhangwl 1 week ago
parent
commit
c77471776b
8 changed files with 87 additions and 49 deletions
  1. 1 1
      mildoc_admin/.env
  2. 7 5
      mildoc_index/.env
  3. 47 39
      mildoc_index/embedding.py
  4. 2 1
      mildoc_index/requirements.txt
  5. 2 2
      mildoc_wxkf/.env
  6. 1 0
      pyproject.toml
  7. 2 1
      requirements-all.txt
  8. 25 0
      uv.lock

+ 1 - 1
mildoc_admin/.env

@@ -27,7 +27,7 @@ MILVUS_PASSWORD=admin123
 MILVUS_DATABASE=mildoc
 MILVUS_DATABASE=mildoc
 MILVUS_COLLECTION=mildoc_collection
 MILVUS_COLLECTION=mildoc_collection
 MILVUS_INDEX_NAME=mildoc_index
 MILVUS_INDEX_NAME=mildoc_index
-MILVUS_VECTOR_DIM=768
+MILVUS_VECTOR_DIM=1024
 
 
 # 管理后台配置
 # 管理后台配置
 ADMIN_USERNAME=admin
 ADMIN_USERNAME=admin

+ 7 - 5
mildoc_index/.env

@@ -32,14 +32,16 @@ MILVUS_PASSWORD=admin123
 MILVUS_DATABASE=mildoc
 MILVUS_DATABASE=mildoc
 MILVUS_COLLECTION=mildoc_collection
 MILVUS_COLLECTION=mildoc_collection
 MILVUS_INDEX_NAME=mildoc_index
 MILVUS_INDEX_NAME=mildoc_index
-MILVUS_VECTOR_DIM=768
+MILVUS_VECTOR_DIM=1024
 
 
 # bailian embedding llm
 # bailian embedding llm
-OPENAI_API_KEY=sk-
-OPENAI_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1
-ENBEDDING_MODEL=text-embedding-v4
+OPENAI_API_KEY=ark-95d3c2b6-b694-46b8-a3b2-5b573e1596ca-cf7e4
+#OPENAI_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1
+OPENAI_BASE_URL=https://ark.cn-beijing.volces.com/api/v3
+#ENBEDDING_MODEL=text-embedding-v4
+ENBEDDING_MODEL=doubao-embedding-vision-251215  #支持多模态的向量模型
 
 
 
 
 # MinerU
 # MinerU
 # 启用MinerU前,先安装: uv add mineru[core]==2.7.0
 # 启用MinerU前,先安装: uv add mineru[core]==2.7.0
-# MINERU_SERVER_URL=http://127.0.0.1:8000
+# MINERU_SERVER_URL=http://127.0.0.1:8000

+ 47 - 39
mildoc_index/embedding.py

@@ -3,6 +3,7 @@ from openai import OpenAI
 from dotenv import load_dotenv
 from dotenv import load_dotenv
 from typing import List
 from typing import List
 from logger.logging import setup_logging
 from logger.logging import setup_logging
+from volcenginesdkarkruntime import Ark
 
 
 load_dotenv()
 load_dotenv()
 
 
@@ -13,10 +14,17 @@ class EmbeddingTool:
         """
         """
         初始化embedding工具
         初始化embedding工具
         """
         """
-        self.client = OpenAI(
+        # self.client = OpenAI(
+        #     api_key=os.getenv("OPENAI_API_KEY"),
+        #     base_url=os.getenv("OPENAI_BASE_URL")
+        # )
+
+        self.client = Ark(
+            # 从环境变量中获取您的 API Key。此为默认方式,您可根据需要进行修改
             api_key=os.getenv("OPENAI_API_KEY"),
             api_key=os.getenv("OPENAI_API_KEY"),
             base_url=os.getenv("OPENAI_BASE_URL")
             base_url=os.getenv("OPENAI_BASE_URL")
         )
         )
+
         self.model = os.getenv("ENBEDDING_MODEL")
         self.model = os.getenv("ENBEDDING_MODEL")
         self.dimensions = int(os.getenv("MILVUS_VECTOR_DIM"))
         self.dimensions = int(os.getenv("MILVUS_VECTOR_DIM"))
         self.encoding_format = "float"
         self.encoding_format = "float"
@@ -24,68 +32,68 @@ class EmbeddingTool:
         if not self.client or not self.model or not self.dimensions or not self.encoding_format:
         if not self.client or not self.model or not self.dimensions or not self.encoding_format:
             logger.error("Embedding工具初始化失败")
             logger.error("Embedding工具初始化失败")
             raise ValueError("Embedding工具初始化失败")
             raise ValueError("Embedding工具初始化失败")
-        
+
     def get_embedding(self, text: str) -> List[float]:
     def get_embedding(self, text: str) -> List[float]:
         """
         """
         获取单个文本的embedding向量
         获取单个文本的embedding向量
-        
+
         Args:
         Args:
             text (str): 输入文本
             text (str): 输入文本
-            
+
         Returns:
         Returns:
             List[float]: 向量列表
             List[float]: 向量列表
         """
         """
         try:
         try:
-            response = self.client.embeddings.create(
+            response = self.client.multimodal_embeddings.create(
                 model=self.model,
                 model=self.model,
-                input=[text],
+                input=[{"type":"text","text":text}],
                 dimensions=self.dimensions,
                 dimensions=self.dimensions,
                 encoding_format=self.encoding_format
                 encoding_format=self.encoding_format
             )
             )
-            
+
             if not response.data:
             if not response.data:
                 logger.error("获取embedding失败")
                 logger.error("获取embedding失败")
                 return []
                 return []
-            
-            embedding = response.data[0].embedding
+
+            embedding = response.data.embedding
             if not embedding:
             if not embedding:
                 logger.error("获取embedding失败")
                 logger.error("获取embedding失败")
                 return []
                 return []
-            
+
             if len(embedding) != self.dimensions:
             if len(embedding) != self.dimensions:
-                logger.error("获取embedding失败")
+                logger.error("获取dimensions失败")
                 return []
                 return []
-            
+
             return embedding
             return embedding
-            
+
         except Exception as e:
         except Exception as e:
             print(f"获取embedding失败: {e}")
             print(f"获取embedding失败: {e}")
             return []
             return []
-    
+
     def get_embeddings_batch(self, texts: List[str]) -> List[List[float]]:
     def get_embeddings_batch(self, texts: List[str]) -> List[List[float]]:
         """
         """
         批量获取多个文本的embedding向量
         批量获取多个文本的embedding向量
-        
+
         Args:
         Args:
             texts (List[str]): 文本列表
             texts (List[str]): 文本列表
-            
+
         Returns:
         Returns:
             List[List[float]]: 向量列表的列表
             List[List[float]]: 向量列表的列表
         """
         """
         try:
         try:
-            response = self.client.embeddings.create(
+            response = self.client.multimodal_embeddings.create(
                 model=self.model,
                 model=self.model,
                 input=texts,
                 input=texts,
                 dimensions=self.dimensions,
                 dimensions=self.dimensions,
                 encoding_format=self.encoding_format
                 encoding_format=self.encoding_format
             )
             )
-            
+
             if not response.data:
             if not response.data:
                 logger.error("批量获取embedding失败")
                 logger.error("批量获取embedding失败")
                 return []
                 return []
-            
-            embeddings = [data.embedding for data in response.data] 
-            
+
+            embeddings = [data.embedding for data in response.data]
+
             if len(embeddings) != len(texts):
             if len(embeddings) != len(texts):
                 logger.error("批量获取embedding失败")
                 logger.error("批量获取embedding失败")
                 return []
                 return []
@@ -94,17 +102,17 @@ class EmbeddingTool:
                 if len(embedding) != self.dimensions:
                 if len(embedding) != self.dimensions:
                     logger.error("批量获取embedding失败")
                     logger.error("批量获取embedding失败")
                     return []
                     return []
-            
+
             return embeddings
             return embeddings
-            
+
         except Exception as e:
         except Exception as e:
             print(f"批量获取embedding失败: {e}")
             print(f"批量获取embedding失败: {e}")
             return []
             return []
-    
+
     def get_model_info(self) -> dict:
     def get_model_info(self) -> dict:
         """
         """
         获取模型信息
         获取模型信息
-        
+
         Returns:
         Returns:
             dict: 模型信息
             dict: 模型信息
         """
         """
@@ -112,20 +120,20 @@ class EmbeddingTool:
             "model": self.model,
             "model": self.model,
             "dimensions": self.dimensions,
             "dimensions": self.dimensions,
             "encoding_format": self.encoding_format,
             "encoding_format": self.encoding_format,
-            "base_url": self.client.base_url
+            "base_url": self.client._base_url
         }
         }
 
 
 
 
 # 使用示例和测试
 # 使用示例和测试
 if __name__ == "__main__":
 if __name__ == "__main__":
     print("=== Embedding工具测试 ===")
     print("=== Embedding工具测试 ===")
-    
+
     # 创建embedding工具实例
     # 创建embedding工具实例
     embedding_tool = EmbeddingTool()
     embedding_tool = EmbeddingTool()
-    
+
     # 显示模型信息
     # 显示模型信息
     print("模型信息:", embedding_tool.get_model_info())
     print("模型信息:", embedding_tool.get_model_info())
-    
+
     # 单个文本embedding测试
     # 单个文本embedding测试
     print("\n=== 单个文本embedding测试 ===")
     print("\n=== 单个文本embedding测试 ===")
     test_text = "这是一个测试文档的内容,用于生成向量表示。"
     test_text = "这是一个测试文档的内容,用于生成向量表示。"
@@ -134,21 +142,21 @@ if __name__ == "__main__":
         print(f"文本: {test_text}")
         print(f"文本: {test_text}")
         print(f"向量维度: {len(embedding)}")
         print(f"向量维度: {len(embedding)}")
         print(f"向量前5个值: {embedding[:5]}")
         print(f"向量前5个值: {embedding[:5]}")
-    
+
     # 批量文本embedding测试
     # 批量文本embedding测试
     print("\n=== 批量文本embedding测试 ===")
     print("\n=== 批量文本embedding测试 ===")
     test_texts = [
     test_texts = [
         "风急天高猿啸哀",
         "风急天高猿啸哀",
-        "渚清沙白鸟飞回", 
+        "渚清沙白鸟飞回",
         "无边落木萧萧下",
         "无边落木萧萧下",
         "不尽长江滚滚来"
         "不尽长江滚滚来"
     ]
     ]
-    
-    embeddings = embedding_tool.get_embeddings_batch(test_texts)
-    if embeddings:
-        print(f"处理了 {len(embeddings)} 个文本")
-        for i, (text, emb) in enumerate(zip(test_texts, embeddings)):
-            print(f"文本{i+1}: {text}")
-            print(f"向量维度: {len(emb)}, 前3个值: {emb[:5]}")
-    
+
+    # embeddings = embedding_tool.get_embeddings_batch(test_texts)
+    # if embeddings:
+    #     print(f"处理了 {len(embeddings)} 个文本")
+    #     for i, (text, emb) in enumerate(zip(test_texts, embeddings)):
+    #         print(f"文本{i+1}: {text}")
+    #         print(f"向量维度: {len(emb)}, 前3个值: {emb[:5]}")
+
     print("\n=== 测试完成 ===")
     print("\n=== 测试完成 ===")

+ 2 - 1
mildoc_index/requirements.txt

@@ -6,6 +6,7 @@ pymilvus==2.6.5
 
 
 # OpenAI API客户端(用于embedding)
 # OpenAI API客户端(用于embedding)
 openai==2.11.0
 openai==2.11.0
+volcengine-python-sdk[ark]==5.0.17
 
 
 # 环境变量管理
 # 环境变量管理
 python-dotenv==1.2.1
 python-dotenv==1.2.1
@@ -23,4 +24,4 @@ langchain-text-splitters==1.1.0
 
 
 # 阿里云
 # 阿里云
 alibabacloud-oss-v2==1.2.1
 alibabacloud-oss-v2==1.2.1
-aliyun-mns-sdk==1.3.0
+aliyun-mns-sdk==1.3.0

+ 2 - 2
mildoc_wxkf/.env

@@ -22,7 +22,7 @@ MILVUS_PASSWORD=admin123
 MILVUS_DATABASE=mildoc
 MILVUS_DATABASE=mildoc
 MILVUS_COLLECTION=mildoc_collection
 MILVUS_COLLECTION=mildoc_collection
 MILVUS_INDEX_NAME=mildoc_index
 MILVUS_INDEX_NAME=mildoc_index
-MILVUS_VECTOR_DIM=768
+MILVUS_VECTOR_DIM=1024
 
 
 # LLM配置 百炼
 # LLM配置 百炼
 LLM_MODEL_NAME=qwen-plus
 LLM_MODEL_NAME=qwen-plus
@@ -46,4 +46,4 @@ RERANK_ENDPOINT=https://dashscope.aliyuncs.com/api/v1/services/rerank/text-reran
 # LANGFUSE_ENABLE=false
 # LANGFUSE_ENABLE=false
 # LANGFUSE_SECRET_KEY=sk-lf-3badb70b-3482-4f87-9581-7b2118ea43fc
 # LANGFUSE_SECRET_KEY=sk-lf-3badb70b-3482-4f87-9581-7b2118ea43fc
 # LANGFUSE_PUBLIC_KEY=pk-lf-9c98917b-7875-4f34-8f72-1afd87491f84
 # LANGFUSE_PUBLIC_KEY=pk-lf-9c98917b-7875-4f34-8f72-1afd87491f84
-# LANGFUSE_BASE_URL=https://cdbece33249643e581c6525392ba1dda--3000.ap-shanghai2.cloudstudio.club
+# LANGFUSE_BASE_URL=https://cdbece33249643e581c6525392ba1dda--3000.ap-shanghai2.cloudstudio.club

+ 1 - 0
pyproject.toml

@@ -22,6 +22,7 @@ dependencies = [
     # "mineru[core]==2.7.0",
     # "mineru[core]==2.7.0",
     "minio==7.2.20",
     "minio==7.2.20",
     "openai==2.11.0",
     "openai==2.11.0",
+    "volcengine-python-sdk[ark]==5.0.42",
     "pycryptodome==3.23.0",
     "pycryptodome==3.23.0",
     "pymilvus==2.6.5",
     "pymilvus==2.6.5",
     "pypdf2==3.0.1",
     "pypdf2==3.0.1",

+ 2 - 1
requirements-all.txt

@@ -26,6 +26,7 @@ pymilvus==2.6.5
 
 
 # OpenAI API客户端(用于embedding)
 # OpenAI API客户端(用于embedding)
 openai==2.11.0
 openai==2.11.0
+volcengine-python-sdk[ark]==5.0.17
 
 
 # langchain相关依赖
 # langchain相关依赖
 langchain==1.1.3
 langchain==1.1.3
@@ -37,4 +38,4 @@ langchain-text-splitters==1.1.0
 
 
 # 阿里云
 # 阿里云
 alibabacloud-oss-v2==1.2.1
 alibabacloud-oss-v2==1.2.1
-aliyun-mns-sdk==1.3.0
+aliyun-mns-sdk==1.3.0

+ 25 - 0
uv.lock

@@ -1049,6 +1049,7 @@ dependencies = [
     { name = "python-dotenv" },
     { name = "python-dotenv" },
     { name = "pytz" },
     { name = "pytz" },
     { name = "requests" },
     { name = "requests" },
+    { name = "volcengine-python-sdk", extra = ["ark"] },
 ]
 ]
 
 
 [package.metadata]
 [package.metadata]
@@ -1075,6 +1076,7 @@ requires-dist = [
     { name = "python-dotenv", specifier = "==1.2.1" },
     { name = "python-dotenv", specifier = "==1.2.1" },
     { name = "pytz", specifier = "==2025.2" },
     { name = "pytz", specifier = "==2025.2" },
     { name = "requests", specifier = "==2.32.5" },
     { name = "requests", specifier = "==2.32.5" },
+    { name = "volcengine-python-sdk", extras = ["ark"], specifier = "==5.0.42" },
 ]
 ]
 
 
 [[package]]
 [[package]]
@@ -1841,6 +1843,29 @@ wheels = [
 ]
 ]
 
 
 [[package]]
 [[package]]
+name = "volcengine-python-sdk"
+version = "5.0.42"
+source = { registry = "https://mirrors.aliyun.com/pypi/simple/" }
+dependencies = [
+    { name = "certifi" },
+    { name = "python-dateutil" },
+    { name = "six" },
+    { name = "urllib3" },
+]
+sdist = { url = "https://mirrors.aliyun.com/pypi/packages/47/6d/121fc4d09ae5da241c9771ddc458516c28f8ad8fa3059910863a9a5953b3/volcengine_python_sdk-5.0.42.tar.gz", hash = "sha256:93d735a3406425c4438625696fa7238130e0e3a5ace6c7056d79275a3595d034" }
+wheels = [
+    { url = "https://mirrors.aliyun.com/pypi/packages/d3/c7/0ed066988d558dbf0942bb08e750e5939fbd934436a540a29497b210d6cc/volcengine_python_sdk-5.0.42-py2.py3-none-any.whl", hash = "sha256:899bbb329f839d76cd5bb4c39dfc5fd53e62677424c2f4a119b8e52b67b6f8a7" },
+]
+
+[package.optional-dependencies]
+ark = [
+    { name = "anyio" },
+    { name = "cryptography" },
+    { name = "httpx" },
+    { name = "pydantic" },
+]
+
+[[package]]
 name = "werkzeug"
 name = "werkzeug"
 version = "3.1.5"
 version = "3.1.5"
 source = { registry = "https://mirrors.aliyun.com/pypi/simple/" }
 source = { registry = "https://mirrors.aliyun.com/pypi/simple/" }