#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 基于FAISS的多模态检索系统 支持文搜文、文搜图、图搜文、图搜图四种检索模式 """ import torch import torch.nn as nn from torch.nn.parallel import DataParallel import numpy as np from PIL import Image from transformers import AutoModel, AutoProcessor, AutoTokenizer from typing import List, Union, Tuple, Dict, Any, Optional import os import json from pathlib import Path import logging import gc from concurrent.futures import ThreadPoolExecutor, as_completed import threading from faiss_vector_store import FaissVectorStore # 设置日志 logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class MultimodalRetrievalFAISS: """基于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: int = 12, index_path: str = "faiss_index"): """ 初始化多模态检索系统 Args: model_name: 模型名称 use_all_gpus: 是否使用所有可用GPU gpu_ids: 指定使用的GPU ID列表 min_memory_gb: 最小可用内存(GB) index_path: FAISS索引文件路径 """ self.model_name = model_name self.index_path = index_path # 设置GPU设备 self._setup_devices(use_all_gpus, gpu_ids, min_memory_gb) # 清理GPU内存 self._clear_all_gpu_memory() # 加载模型和处理器 self._load_model_and_processor() # 初始化FAISS向量存储 self.vector_store = FaissVectorStore( index_path=index_path, dimension=3584 # OpenSearch-AI/Ops-MM-embedding-v1-7B的向量维度 ) logger.info(f"多模态检索系统初始化完成，使用模型: {model_name}") logger.info(f"向量存储路径: {index_path}") def _setup_devices(self, use_all_gpus: bool, gpu_ids: List[int], min_memory_gb: int): """设置GPU设备""" self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.use_gpu = self.device.type == "cuda" if self.use_gpu: self.available_gpus = self._get_available_gpus(min_memory_gb) if not self.available_gpus: logger.warning(f"没有可用的GPU或GPU内存不足{min_memory_gb}GB，将使用CPU") self.device = torch.device("cpu") self.use_gpu = False else: if gpu_ids: self.gpu_ids = [gid for gid in gpu_ids if gid in self.available_gpus] if not self.gpu_ids: logger.warning(f"指定的GPU {gpu_ids}不可用或内存不足，将使用可用的GPU: {self.available_gpus}") self.gpu_ids = self.available_gpus elif use_all_gpus: self.gpu_ids = self.available_gpus else: self.gpu_ids = [self.available_gpus[0]] logger.info(f"使用GPU: {self.gpu_ids}") self.device = torch.device(f"cuda:{self.gpu_ids[0]}") def _get_available_gpus(self, min_memory_gb: int) -> List[int]: """获取可用的GPU列表""" available_gpus = [] for i in range(torch.cuda.device_count()): total_mem = torch.cuda.get_device_properties(i).total_memory / (1024 ** 3) # GB if total_mem >= min_memory_gb: available_gpus.append(i) return available_gpus def _clear_all_gpu_memory(self): """清理GPU内存""" if torch.cuda.is_available(): torch.cuda.empty_cache() gc.collect() def _load_model_and_processor(self): """加载模型和处理器""" logger.info(f"加载模型和处理器: {self.model_name}") # 加载tokenizer和processor self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) self.processor = AutoProcessor.from_pretrained(self.model_name) # 加载模型 self.model = AutoModel.from_pretrained( self.model_name, torch_dtype=torch.float16 if self.use_gpu else torch.float32, device_map="auto" if len(self.gpu_ids) > 1 else None ) # 如果使用多GPU，包装模型 if len(self.gpu_ids) > 1: self.model = DataParallel(self.model, device_ids=self.gpu_ids) self.model.eval() self.model.to(self.device) logger.info("模型和处理器加载完成") def encode_text(self, text: Union[str, List[str]]) -> np.ndarray: """编码文本为向量""" if isinstance(text, str): text = [text] inputs = self.tokenizer(text, padding=True, truncation=True, return_tensors="pt") inputs = {k: v.to(self.device) for k, v in inputs.items()} with torch.no_grad(): outputs = self.model(**inputs) # 获取[CLS]标记的隐藏状态作为句子表示 text_embeddings = outputs.last_hidden_state[:, 0, :].cpu().numpy() # 归一化向量 text_embeddings = text_embeddings / np.linalg.norm(text_embeddings, axis=1, keepdims=True) return text_embeddings[0] if len(text) == 1 else text_embeddings def encode_image(self, image: Union[Image.Image, List[Image.Image]]) -> np.ndarray: """编码图像为向量""" if isinstance(image, Image.Image): image = [image] inputs = self.processor(images=image, return_tensors="pt") inputs = {k: v.to(self.device) for k, v in inputs.items()} with torch.no_grad(): outputs = self.model.vision_model(**inputs) # 获取[CLS]标记的隐藏状态作为图像表示 image_embeddings = outputs.pooler_output.cpu().numpy() # 归一化向量 image_embeddings = image_embeddings / np.linalg.norm(image_embeddings, axis=1, keepdims=True) return image_embeddings[0] if len(image) == 1 else image_embeddings def add_texts( self, texts: List[str], metadatas: Optional[List[Dict[str, Any]]] = None ) -> List[str]: """ 添加文本到检索系统 Args: texts: 文本列表 metadatas: 元数据列表，每个元素是一个字典 Returns: 添加的文本ID列表 """ if not texts: return [] if metadatas is None: metadatas = [{} for _ in range(len(texts))] if len(texts) != len(metadatas): raise ValueError("texts和metadatas长度必须相同") # 编码文本 text_embeddings = self.encode_text(texts) # 准备元数据 for i, text in enumerate(texts): metadatas[i].update({ "text": text, "type": "text" }) # 添加到向量存储 vector_ids = self.vector_store.add_vectors(text_embeddings, metadatas) logger.info(f"成功添加{len(vector_ids)}条文本到检索系统") return vector_ids def add_images( self, images: List[Image.Image], metadatas: Optional[List[Dict[str, Any]]] = None ) -> List[str]: """ 添加图像到检索系统 Args: images: PIL图像列表 metadatas: 元数据列表，每个元素是一个字典 Returns: 添加的图像ID列表 """ if not images: return [] if metadatas is None: metadatas = [{} for _ in range(len(images))] if len(images) != len(metadatas): raise ValueError("images和metadatas长度必须相同") # 编码图像 image_embeddings = self.encode_image(images) # 准备元数据 for i, image in enumerate(images): metadatas[i].update({ "type": "image", "width": image.width, "height": image.height }) # 添加到向量存储 vector_ids = self.vector_store.add_vectors(image_embeddings, metadatas) logger.info(f"成功添加{len(vector_ids)}张图像到检索系统") return vector_ids def search_by_text( self, query: str, k: int = 5, filter_condition: Optional[Dict[str, Any]] = None ) -> List[Dict[str, Any]]: """ 文本搜索 Args: query: 查询文本 k: 返回结果数量 filter_condition: 过滤条件 Returns: 搜索结果列表，每个元素包含相似项和分数 """ # 编码查询文本 query_embedding = self.encode_text(query) # 执行搜索 results, distances = self.vector_store.search(query_embedding, k) # 处理结果 search_results = [] for i, (result, distance) in enumerate(zip(results, distances)): result["score"] = 1.0 / (1.0 + distance) # 将距离转换为相似度分数 search_results.append(result) return search_results def search_by_image( self, image: Image.Image, k: int = 5, filter_condition: Optional[Dict[str, Any]] = None ) -> List[Dict[str, Any]]: """ 图像搜索 Args: image: 查询图像 k: 返回结果数量 filter_condition: 过滤条件 Returns: 搜索结果列表，每个元素包含相似项和分数 """ # 编码查询图像 query_embedding = self.encode_image(image) # 执行搜索 results, distances = self.vector_store.search(query_embedding, k) # 处理结果 search_results = [] for i, (result, distance) in enumerate(zip(results, distances)): result["score"] = 1.0 / (1.0 + distance) # 将距离转换为相似度分数 search_results.append(result) return search_results def get_vector_count(self) -> int: """获取向量数量""" return self.vector_store.get_vector_count() def save_index(self): """保存索引""" self.vector_store.save_index() logger.info("索引已保存") def __del__(self): """析构函数，确保资源被正确释放""" if hasattr(self, 'model'): del self.model self._clear_all_gpu_memory() if hasattr(self, 'vector_store'): self.save_index() def test_faiss_system(): """测试FAISS多模态检索系统""" import time from PIL import Image import numpy as np # 初始化检索系统 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("\n测试文本搜索...") query_text = "一只猫在睡觉" print(f"查询: {query_text}") results = retrieval.search_by_text(query_text, 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_faiss_system()