444 lines
17 KiB
Python
444 lines
17 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
纯百度VDB多模态检索系统 - 完全替代FAISS
|
||
支持文搜文、文搜图、图搜文、图搜图四种检索模式
|
||
"""
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
from torch.nn.parallel import DataParallel, DistributedDataParallel
|
||
import numpy as np
|
||
from PIL import Image
|
||
from transformers import AutoModel, AutoProcessor, AutoTokenizer
|
||
from typing import List, Union, Tuple, Dict, Any
|
||
import os
|
||
import json
|
||
from pathlib import Path
|
||
import logging
|
||
import gc
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
import threading
|
||
|
||
from baidu_vdb_production import BaiduVDBProduction
|
||
|
||
# 设置日志
|
||
logging.basicConfig(level=logging.INFO)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
class MultimodalRetrievalVDBOnly:
|
||
"""纯百度VDB多模态检索系统，完全替代FAISS"""
|
||
|
||
def __init__(self, model_name: str = "OpenSearch-AI/Ops-MM-embedding-v1-7B",
|
||
use_all_gpus: bool = True, gpu_ids: List[int] = None, min_memory_gb=12):
|
||
"""
|
||
初始化纯VDB多模态检索系统
|
||
|
||
Args:
|
||
model_name: 模型名称
|
||
use_all_gpus: 是否使用所有可用GPU
|
||
gpu_ids: 指定使用的GPU ID列表
|
||
min_memory_gb: 最小可用内存（GB）
|
||
"""
|
||
self.model_name = model_name
|
||
|
||
# 设置GPU设备
|
||
self._setup_devices(use_all_gpus, gpu_ids, min_memory_gb)
|
||
|
||
# 清理GPU内存
|
||
self._clear_all_gpu_memory()
|
||
|
||
logger.info(f"正在加载模型到多GPU: {self.device_ids}")
|
||
|
||
# 加载模型和处理器
|
||
self.model = None
|
||
self.tokenizer = None
|
||
self.processor = None
|
||
self._load_model_multigpu()
|
||
|
||
# 初始化百度VDB后端（替代FAISS索引）
|
||
logger.info("初始化百度VDB后端...")
|
||
self.vdb = BaiduVDBProduction()
|
||
logger.info("✅ 百度VDB后端初始化完成")
|
||
|
||
# 线程锁
|
||
self.model_lock = threading.Lock()
|
||
|
||
logger.info("✅ 纯VDB多模态检索系统初始化完成")
|
||
|
||
def _setup_devices(self, use_all_gpus, gpu_ids, min_memory_gb):
|
||
"""设置GPU设备"""
|
||
if not torch.cuda.is_available():
|
||
raise RuntimeError("CUDA不可用，需要GPU支持")
|
||
|
||
total_gpus = torch.cuda.device_count()
|
||
logger.info(f"检测到 {total_gpus} 个GPU")
|
||
|
||
# 获取可用GPU
|
||
available_gpus = []
|
||
for i in range(total_gpus):
|
||
memory_gb = torch.cuda.get_device_properties(i).total_memory / (1024**3)
|
||
free_memory = torch.cuda.memory_reserved(i) / (1024**3)
|
||
available_memory = memory_gb - free_memory
|
||
|
||
logger.info(f"GPU {i}: {torch.cuda.get_device_properties(i).name} ({memory_gb:.1f}GB)")
|
||
|
||
if available_memory >= min_memory_gb:
|
||
available_gpus.append(i)
|
||
logger.info(f"GPU {i}: {available_memory:.0f}MB 可用 (合适)")
|
||
else:
|
||
logger.info(f"GPU {i}: {available_memory:.0f}MB 可用 (不足)")
|
||
|
||
if not available_gpus:
|
||
raise RuntimeError(f"没有找到满足 {min_memory_gb}GB 内存要求的GPU")
|
||
|
||
# 选择使用的GPU
|
||
if gpu_ids:
|
||
self.device_ids = [gpu_id for gpu_id in gpu_ids if gpu_id in available_gpus]
|
||
elif use_all_gpus:
|
||
self.device_ids = available_gpus
|
||
else:
|
||
self.device_ids = [available_gpus[0]]
|
||
|
||
if not self.device_ids:
|
||
raise RuntimeError("没有可用的GPU设备")
|
||
|
||
# 设置主设备
|
||
self.primary_device = f"cuda:{self.device_ids[0]}"
|
||
torch.cuda.set_device(self.device_ids[0])
|
||
|
||
logger.info(f"使用GPU: {self.device_ids}, 主设备: {self.primary_device}")
|
||
|
||
def _clear_all_gpu_memory(self):
|
||
"""清理所有GPU内存"""
|
||
for device_id in self.device_ids:
|
||
with torch.cuda.device(device_id):
|
||
torch.cuda.empty_cache()
|
||
gc.collect()
|
||
logger.info("所有GPU内存已清理")
|
||
|
||
def _load_model_multigpu(self):
|
||
"""加载模型到多GPU"""
|
||
try:
|
||
# 清理GPU内存
|
||
self._clear_all_gpu_memory()
|
||
|
||
logger.info(f"正在加载模型到多GPU: {self.device_ids}")
|
||
|
||
# 加载模型
|
||
self.model = AutoModel.from_pretrained(
|
||
self.model_name,
|
||
torch_dtype=torch.float16,
|
||
trust_remote_code=True,
|
||
device_map="auto"
|
||
)
|
||
|
||
# 加载tokenizer和processor
|
||
self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, trust_remote_code=True)
|
||
logger.info("Tokenizer加载成功")
|
||
|
||
self.processor = AutoProcessor.from_pretrained(self.model_name, trust_remote_code=True)
|
||
logger.info("Processor加载成功")
|
||
|
||
# 显示设备映射
|
||
if hasattr(self.model, 'hf_device_map'):
|
||
logger.info(f"模型已成功加载到设备: {dict(list(self.model.hf_device_map.items())[:10])}")
|
||
|
||
self.model.eval()
|
||
logger.info("多GPU模型加载完成")
|
||
|
||
except Exception as e:
|
||
logger.error(f"模型加载失败: {e}")
|
||
raise
|
||
|
||
def encode_text_batch(self, texts: List[str], batch_size: int = 8) -> np.ndarray:
|
||
"""批量编码文本"""
|
||
try:
|
||
with self.model_lock:
|
||
all_embeddings = []
|
||
|
||
for i in range(0, len(texts), batch_size):
|
||
batch_texts = texts[i:i + batch_size]
|
||
|
||
# 使用processor处理文本
|
||
inputs = self.processor(
|
||
text=batch_texts,
|
||
return_tensors="pt",
|
||
padding=True,
|
||
truncation=True,
|
||
max_length=512
|
||
)
|
||
|
||
# 将输入移动到主设备
|
||
inputs = {k: v.to(self.primary_device) for k, v in inputs.items()}
|
||
|
||
with torch.no_grad():
|
||
outputs = self.model(**inputs)
|
||
embeddings = outputs.last_hidden_state.mean(dim=1)
|
||
embeddings = embeddings.cpu().numpy()
|
||
all_embeddings.append(embeddings)
|
||
|
||
return np.vstack(all_embeddings)
|
||
|
||
except Exception as e:
|
||
logger.error(f"文本编码失败: {e}")
|
||
return np.zeros((len(texts), 3584), dtype=np.float32)
|
||
|
||
def encode_image_batch(self, images: List[Union[str, Image.Image]], batch_size: int = 4) -> np.ndarray:
|
||
"""批量编码图像"""
|
||
try:
|
||
with self.model_lock:
|
||
processed_images = []
|
||
|
||
# 处理图像输入
|
||
for img in images:
|
||
if isinstance(img, str):
|
||
if os.path.exists(img):
|
||
processed_images.append(Image.open(img).convert('RGB'))
|
||
else:
|
||
logger.warning(f"图像文件不存在: {img}")
|
||
processed_images.append(Image.new('RGB', (224, 224), color='white'))
|
||
elif isinstance(img, Image.Image):
|
||
processed_images.append(img.convert('RGB'))
|
||
else:
|
||
logger.warning(f"不支持的图像类型: {type(img)}")
|
||
processed_images.append(Image.new('RGB', (224, 224), color='white'))
|
||
|
||
all_embeddings = []
|
||
|
||
for i in range(0, len(processed_images), batch_size):
|
||
batch_images = processed_images[i:i + batch_size]
|
||
|
||
# 使用processor处理图像
|
||
inputs = self.processor(
|
||
images=batch_images,
|
||
return_tensors="pt",
|
||
padding=True
|
||
)
|
||
|
||
# 将输入移动到主设备
|
||
inputs = {k: v.to(self.primary_device) for k, v in inputs.items()}
|
||
|
||
with torch.no_grad():
|
||
outputs = self.model(**inputs)
|
||
embeddings = outputs.last_hidden_state.mean(dim=1)
|
||
embeddings = embeddings.cpu().numpy()
|
||
all_embeddings.append(embeddings)
|
||
|
||
return np.vstack(all_embeddings)
|
||
|
||
except Exception as e:
|
||
logger.error(f"图像编码失败: {e}")
|
||
embedding_dim = 3584
|
||
embeddings = np.zeros((len(images), embedding_dim), dtype=np.float32)
|
||
return embeddings
|
||
|
||
def build_text_index_parallel(self, texts: List[str], save_path: str = None):
|
||
"""
|
||
构建文本索引（使用VDB替代FAISS）
|
||
"""
|
||
try:
|
||
logger.info(f"正在构建文本索引，共 {len(texts)} 条文本")
|
||
|
||
# 编码文本
|
||
embeddings = self.encode_text_batch(texts)
|
||
|
||
# 使用VDB存储
|
||
self.vdb.build_text_index(texts, embeddings)
|
||
|
||
logger.info("文本索引构建完成")
|
||
|
||
except Exception as e:
|
||
logger.error(f"构建文本索引失败: {e}")
|
||
raise
|
||
|
||
def build_image_index_parallel(self, image_paths: List[str], save_path: str = None):
|
||
"""
|
||
构建图像索引（使用VDB替代FAISS）
|
||
"""
|
||
try:
|
||
logger.info(f"正在构建图像索引，共 {len(image_paths)} 张图像")
|
||
|
||
# 编码图像
|
||
embeddings = self.encode_image_batch(image_paths)
|
||
|
||
# 使用VDB存储
|
||
self.vdb.build_image_index(image_paths, embeddings)
|
||
|
||
logger.info("图像索引构建完成")
|
||
|
||
except Exception as e:
|
||
logger.error(f"构建图像索引失败: {e}")
|
||
raise
|
||
|
||
def search_text_by_text(self, query: str, top_k: int = 5) -> List[Tuple[str, float]]:
|
||
"""文搜文：使用文本查询搜索相似文本"""
|
||
try:
|
||
query_embedding = self.encode_text_batch([query])
|
||
return self.vdb.search_text_by_text(query_embedding[0], top_k)
|
||
except Exception as e:
|
||
logger.error(f"文搜文失败: {e}")
|
||
return []
|
||
|
||
def search_images_by_text(self, query: str, top_k: int = 5) -> List[Tuple[str, float]]:
|
||
"""文搜图：使用文本查询搜索相似图像"""
|
||
try:
|
||
query_embedding = self.encode_text_batch([query])
|
||
return self.vdb.search_images_by_text(query_embedding[0], top_k)
|
||
except Exception as e:
|
||
logger.error(f"文搜图失败: {e}")
|
||
return []
|
||
|
||
def search_images_by_image(self, query_image: Union[str, Image.Image], top_k: int = 5) -> List[Tuple[str, float]]:
|
||
"""图搜图：使用图像查询搜索相似图像"""
|
||
try:
|
||
query_embedding = self.encode_image_batch([query_image])
|
||
return self.vdb.search_images_by_image(query_embedding[0], top_k)
|
||
except Exception as e:
|
||
logger.error(f"图搜图失败: {e}")
|
||
return []
|
||
|
||
def search_text_by_image(self, query_image: Union[str, Image.Image], top_k: int = 5) -> List[Tuple[str, float]]:
|
||
"""图搜文：使用图像查询搜索相似文本"""
|
||
try:
|
||
query_embedding = self.encode_image_batch([query_image])
|
||
return self.vdb.search_text_by_image(query_embedding[0], top_k)
|
||
except Exception as e:
|
||
logger.error(f"图搜文失败: {e}")
|
||
return []
|
||
|
||
# Web应用兼容方法
|
||
def search_by_text(self, query: str, top_k: int = 5) -> List[Tuple[str, float]]:
|
||
"""文搜文：Web应用兼容方法"""
|
||
return self.search_text_by_text(query, top_k)
|
||
|
||
def search_by_image(self, query_image: Union[str, Image.Image], top_k: int = 5) -> List[Tuple[str, float]]:
|
||
"""图搜图：Web应用兼容方法"""
|
||
return self.search_images_by_image(query_image, top_k)
|
||
|
||
def search_images_by_text_query(self, query: str, top_k: int = 5) -> List[Tuple[str, float]]:
|
||
"""文搜图：Web应用兼容方法"""
|
||
return self.search_images_by_text(query, top_k)
|
||
|
||
def search_texts_by_image_query(self, query_image: Union[str, Image.Image], top_k: int = 5) -> List[Tuple[str, float]]:
|
||
"""图搜文：Web应用兼容方法"""
|
||
return self.search_text_by_image(query_image, top_k)
|
||
|
||
def get_statistics(self) -> Dict[str, Any]:
|
||
"""获取系统统计信息"""
|
||
try:
|
||
vdb_stats = self.vdb.get_statistics()
|
||
|
||
stats = {
|
||
"model_name": self.model_name,
|
||
"device_ids": self.device_ids,
|
||
"primary_device": self.primary_device,
|
||
"backend": "Baidu VDB (No FAISS)",
|
||
**vdb_stats
|
||
}
|
||
|
||
return stats
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取统计信息失败: {e}")
|
||
return {"status": "error", "error": str(e)}
|
||
|
||
def clear_all_data(self):
|
||
"""清空所有数据"""
|
||
try:
|
||
self.vdb.clear_all_data()
|
||
logger.info("✅ 所有数据已清空")
|
||
except Exception as e:
|
||
logger.error(f"❌ 清空数据失败: {e}")
|
||
|
||
def get_gpu_memory_info(self):
|
||
"""获取所有GPU内存使用信息"""
|
||
memory_info = {}
|
||
for device_id in self.device_ids:
|
||
with torch.cuda.device(device_id):
|
||
allocated = torch.cuda.memory_allocated() / (1024**3)
|
||
reserved = torch.cuda.memory_reserved() / (1024**3)
|
||
total = torch.cuda.get_device_properties(device_id).total_memory / (1024**3)
|
||
|
||
memory_info[f"GPU_{device_id}"] = {
|
||
"allocated_GB": round(allocated, 2),
|
||
"reserved_GB": round(reserved, 2),
|
||
"total_GB": round(total, 2),
|
||
"free_GB": round(total - reserved, 2)
|
||
}
|
||
|
||
return memory_info
|
||
|
||
def cleanup(self):
|
||
"""清理资源"""
|
||
try:
|
||
if self.vdb:
|
||
self.vdb.close()
|
||
|
||
self._clear_all_gpu_memory()
|
||
logger.info("✅ 资源清理完成")
|
||
except Exception as e:
|
||
logger.error(f"❌ 资源清理失败: {e}")
|
||
|
||
def test_vdb_only_system():
|
||
"""测试纯VDB多模态检索系统"""
|
||
print("=" * 60)
|
||
print("测试纯百度VDB多模态检索系统")
|
||
print("=" * 60)
|
||
|
||
system = None
|
||
|
||
try:
|
||
# 1. 初始化系统
|
||
print("1. 初始化纯VDB多模态检索系统...")
|
||
system = MultimodalRetrievalVDBOnly()
|
||
print("✅ 系统初始化成功")
|
||
|
||
# 2. 构建文本索引
|
||
print("\n2. 构建文本索引...")
|
||
test_texts = [
|
||
"人工智能技术的发展趋势",
|
||
"机器学习在医疗领域的应用",
|
||
"深度学习算法优化方法",
|
||
"计算机视觉技术创新",
|
||
"自然语言处理最新进展"
|
||
]
|
||
|
||
system.build_text_index_parallel(test_texts)
|
||
print("✅ 文本索引构建完成")
|
||
|
||
# 3. 测试文搜文
|
||
print("\n3. 测试文搜文...")
|
||
query = "AI技术"
|
||
results = system.search_text_by_text(query, top_k=3)
|
||
print(f"查询: {query}")
|
||
for i, (text, score) in enumerate(results, 1):
|
||
print(f" {i}. {text} (相似度: {score:.3f})")
|
||
|
||
# 4. 获取统计信息
|
||
print("\n4. 获取统计信息...")
|
||
stats = system.get_statistics()
|
||
print("系统统计:")
|
||
for key, value in stats.items():
|
||
print(f" {key}: {value}")
|
||
|
||
print(f"\n🎉 纯VDB系统测试完成！")
|
||
print("✅ 完全移除FAISS依赖")
|
||
print("✅ 使用百度VDB作为向量数据库")
|
||
print("✅ 支持多模态检索功能")
|
||
|
||
return True
|
||
|
||
except Exception as e:
|
||
print(f"❌ 测试失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
return False
|
||
|
||
finally:
|
||
if system:
|
||
system.cleanup()
|
||
|
||
if __name__ == "__main__":
|
||
test_vdb_only_system()
|