59 lines
1.7 KiB
Python
59 lines
1.7 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
FAISS多模态检索系统简单测试
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
from multimodal_retrieval_faiss import MultimodalRetrievalFAISS
|
|
from PIL import Image
|
|
import numpy as np
|
|
|
|
def test_text_retrieval():
|
|
print("=== 测试文本检索 ===")
|
|
|
|
# 初始化检索系统
|
|
print("初始化检索系统...")
|
|
retrieval = MultimodalRetrievalFAISS(
|
|
model_name="OpenSearch-AI/Ops-MM-embedding-v1-7B",
|
|
use_all_gpus=True,
|
|
index_path="faiss_index_test"
|
|
)
|
|
|
|
# 测试文本
|
|
texts = [
|
|
"一只可爱的橘色猫咪在沙发上睡觉",
|
|
"城市夜景中的高楼大厦和车流",
|
|
"阳光明媚的海滩上，人们在冲浪和晒太阳",
|
|
"美味的意大利面配红酒和沙拉",
|
|
"雪山上滑雪的运动员"
|
|
]
|
|
|
|
# 添加文本
|
|
print("\n添加文本到检索系统...")
|
|
text_ids = retrieval.add_texts(texts)
|
|
print(f"添加了{len(text_ids)}条文本")
|
|
print(f"当前向量数量: {retrieval.get_vector_count()}")
|
|
|
|
# 测试文本搜索
|
|
print("\n测试文本搜索...")
|
|
queries = ["一只猫在睡觉", "都市风光", "海边的景色"]
|
|
|
|
for query in queries:
|
|
print(f"\n查询: {query}")
|
|
results = retrieval.search_by_text(query, k=2)
|
|
for i, result in enumerate(results):
|
|
print(f" 结果 {i+1}: {result.get('text', 'N/A')} (分数: {result.get('score', 0):.4f})")
|
|
|
|
# 保存索引
|
|
print("\n保存索引...")
|
|
retrieval.save_index()
|
|
|
|
print("\n测试完成！")
|
|
|
|
if __name__ == "__main__":
|
|
test_text_retrieval()
|