|
@@ -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=== 测试完成 ===")
|