111 lines
4.1 KiB
Python
111 lines
4.1 KiB
Python
|
import requests
|
|||
|
|
import time
|
|||
|
|
|
|||
|
|
# Ollama配置
|
|||
|
|
OLLAMA_HOST = "http://localhost:11434" # 默认Ollama地址
|
|||
|
|
MODEL_NAME = "bge-m3" # 使用的embedding模型
|
|||
|
|
TEXT_TO_EMBED = "测试文本"
|
|||
|
|
|
|||
|
|
# 定义接口URL和对应的请求体结构
|
|||
|
|
ENDPOINTS = {
|
|||
|
|
"api/embeddings": {
|
|||
|
|
"url": f"{OLLAMA_HOST}/api/embeddings", # 原生API路径
|
|||
|
|
"payload": {"model": MODEL_NAME, "prompt": TEXT_TO_EMBED}, # 原生API用prompt字段
|
|||
|
|
},
|
|||
|
|
"v1/embeddings": {
|
|||
|
|
"url": f"{OLLAMA_HOST}/v1/embeddings", # OpenAI兼容API路径
|
|||
|
|
"payload": {"model": MODEL_NAME, "input": TEXT_TO_EMBED}, # OpenAI兼容API用input字段
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
headers = {"Content-Type": "application/json"}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_endpoint(endpoint_name, endpoint_info):
|
|||
|
|
"""测试单个端点并返回结果"""
|
|||
|
|
print(f"\n测试接口: {endpoint_name}")
|
|||
|
|
url = endpoint_info["url"]
|
|||
|
|
payload = endpoint_info["payload"]
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
start_time = time.time()
|
|||
|
|
response = requests.post(url, headers=headers, json=payload)
|
|||
|
|
response_time = time.time() - start_time
|
|||
|
|
|
|||
|
|
print(f"状态码: {response.status_code}")
|
|||
|
|
print(f"响应时间: {response_time:.3f}秒")
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
data = response.json()
|
|||
|
|
|
|||
|
|
# 处理不同接口的响应结构差异
|
|||
|
|
embedding = None
|
|||
|
|
if endpoint_name == "api/embeddings":
|
|||
|
|
embedding = data.get("embedding") # 原生API返回embedding字段
|
|||
|
|
elif endpoint_name == "v1/embeddings":
|
|||
|
|
embedding = data.get("data", [{}])[0].get("embedding") # OpenAI兼容API返回data数组中的embedding
|
|||
|
|
|
|||
|
|
if embedding:
|
|||
|
|
print(f"Embedding向量长度: {len(embedding)}")
|
|||
|
|
return {
|
|||
|
|
"endpoint": endpoint_name,
|
|||
|
|
"status_code": response.status_code,
|
|||
|
|
"response_time": response_time,
|
|||
|
|
"embedding_length": len(embedding),
|
|||
|
|
"embedding": embedding[:5],
|
|||
|
|
}
|
|||
|
|
else:
|
|||
|
|
print("响应中未找到'embedding'字段")
|
|||
|
|
return {"endpoint": endpoint_name, "status_code": response.status_code, "error": "No embedding field in response"}
|
|||
|
|
|
|||
|
|
except ValueError:
|
|||
|
|
print("响应不是有效的JSON格式")
|
|||
|
|
return {"endpoint": endpoint_name, "status_code": response.status_code, "error": "Invalid JSON response"}
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"请求失败: {str(e)}")
|
|||
|
|
return {"endpoint": endpoint_name, "error": str(e)}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def compare_endpoints():
|
|||
|
|
"""比较两个端点的性能"""
|
|||
|
|
results = []
|
|||
|
|
|
|||
|
|
print("=" * 50)
|
|||
|
|
print(f"开始比较Ollama的embeddings接口，使用模型: {MODEL_NAME}")
|
|||
|
|
print("=" * 50)
|
|||
|
|
|
|||
|
|
for endpoint_name, endpoint_info in ENDPOINTS.items():
|
|||
|
|
results.append(test_endpoint(endpoint_name, endpoint_info))
|
|||
|
|
|
|||
|
|
print("\n" + "=" * 50)
|
|||
|
|
print("比较结果摘要:")
|
|||
|
|
print("=" * 50)
|
|||
|
|
|
|||
|
|
successful_results = [res for res in results if "embedding_length" in res]
|
|||
|
|
|
|||
|
|
if len(successful_results) == 2:
|
|||
|
|
if successful_results[0]["embedding_length"] == successful_results[1]["embedding_length"]:
|
|||
|
|
print(f"两个接口返回的embedding维度相同: {successful_results[0]['embedding_length']}")
|
|||
|
|
else:
|
|||
|
|
print("两个接口返回的embedding维度不同:")
|
|||
|
|
for result in successful_results:
|
|||
|
|
print(f"- {result['endpoint']}: {result['embedding_length']}")
|
|||
|
|
|
|||
|
|
print("\nEmbedding前5个元素示例:")
|
|||
|
|
for result in successful_results:
|
|||
|
|
print(f"- {result['endpoint']}: {result['embedding']}")
|
|||
|
|
|
|||
|
|
faster = min(successful_results, key=lambda x: x["response_time"])
|
|||
|
|
slower = max(successful_results, key=lambda x: x["response_time"])
|
|||
|
|
print(f"\n更快的接口: {faster['endpoint']} ({faster['response_time']:.3f}秒 vs {slower['response_time']:.3f}秒)")
|
|||
|
|
else:
|
|||
|
|
print("至少有一个接口未返回有效的embedding数据")
|
|||
|
|
for result in results:
|
|||
|
|
if "error" in result:
|
|||
|
|
print(f"- {result['endpoint']} 错误: {result['error']}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
compare_endpoints()
|