import os import logging from abc import ABC, abstractmethod from typing import List, Dict, Any from sentence_transformers import SentenceTransformer from pymilvus import MilvusClient, DataType import time import gradio as gr # 配置日志 logging.basicConfig( level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s' ) logger = logging.getLogger(__name__) models = [ 'shibing624/text2vec-base-chinese', 'BAAI/bge-small-zh', 'BAAI/bge-base-zh', 'sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2', 'all-MiniLM-L6-v2', 'all-MiniLM-L12-v2', 'multi-qa-mpnet-base-dot-v1', # 'bge-small-en-v1.5', 不兼容 'all-mpnet-base-v2', 'jinaai/jina-embeddings-v3', ] searchers = {} class BaseEmbeddingModel(ABC): @abstractmethod def encode(self, text: str) -> List[float]: pass @property @abstractmethod def dimension(self) -> int: pass @property @abstractmethod def model_name(self) -> str: pass class SentenceTransformerModel(BaseEmbeddingModel): def __init__(self, model_name: str): self.model = SentenceTransformer(model_name, trust_remote_code=True) self._model_name = model_name def encode(self, text: str) -> List[float]: result = self.model.encode(text).tolist() return result @property def dimension(self) -> int: return self.model.get_sentence_embedding_dimension() @property def model_name(self) -> str: return self._model_name class StickerSearcher: def __init__(self, model: BaseEmbeddingModel): self.model = model self.client = MilvusClient(uri='./sticker.db') self.collection_name = f'test_{model.model_name.replace("/", "_").replace("-", "_")}' def init_collection(self) -> bool: try: self.client.drop_collection(collection_name=self.collection_name) self.client.create_collection( collection_name=self.collection_name, dimension=self.model.dimension, primary_field_name='id', auto_id=True ) self.client.create_index( collection_name=self.collection_name, index_type='IVF_SQ8', metric_type='COSINE', params={'nlist': 128}, index_params={} ) self.client.load_collection(self.collection_name) logger.info(f'Collection initialized: {self.collection_name}') return True except Exception as e: logger.error(f'Collection init failed: {str(e)}') return False def store_vector(self, title: str, description: str, tags: List[str], file_path: str): vector = self.model.encode(description) data = [{ 'vector': vector, 'title': title, 'description': description, 'tags': tags, 'file_name': file_path }] self.client.insert(self.collection_name, data) def search(self, query: str, limit: int = 5) -> List[Dict[str, Any]]: start_time = time.time() query_vector = self.model.encode(query) encode_time = time.time() - start_time start_search_time = time.time() results = self.client.search( collection_name=self.collection_name, data=[query_vector], limit=limit, output_fields=['title', 'description', 'tags', 'file_name'] ) search_time = time.time() - start_search_time total_time = encode_time + search_time logger.info(f'模型 {self.model.model_name} Encoding耗时: ${encode_time:.4f}，搜索耗时: {search_time:.4f} 秒, 总耗时: {total_time:.4f} 秒') return results[0] def create_gradio_ui(): async def search_model(model_name: str, query: str): try: if model_name in searchers: return searchers[model_name].search(query) logger.error(f'Model not loaded: {model_name}') return [] except Exception as e: logger.error(f'Search failed: {model_name} | Error: {str(e)}') return [] async def search_all_models(query): if not query: return [] print(f'>>>> Searching From Models {query}') results = [] for model_name in models: result = await search_model(model_name, query) results.append(result) formatted_results = [] max_results = max(len(r) for r in results) for i in range(max_results): row = [i + 1] for model_results in results: if i < len(model_results): result = model_results[i] image_url = f'https://huggingface.co/datasets/Nekoko/StickerSet/resolve/main/{result["entity"]["file_name"]}' row.append(f'![Sticker]({image_url})\n相似度: {result["distance"]:.4f}') else: row.append('-') formatted_results.append(row) return formatted_results def init_collections(): try: client = MilvusClient(uri='./sticker.db') stickers = client.query( collection_name='stickers', filter='', limit=1000, output_fields=['title', 'description', 'tags', 'file_name'] ) logger.info(f'Stickers loaded: {len(stickers)}') def init_model(model_name): try: searcher = StickerSearcher(SentenceTransformerModel(model_name)) if searcher.init_collection(): searchers[model_name] = searcher for sticker in stickers: searcher.store_vector( sticker.get('title'), sticker.get('description'), sticker.get('tags'), sticker.get('file_name') ) logger.info(f'Model initialized: {model_name}') except Exception as e: logger.error(f'Model init failed: {model_name} | Error: {str(e)}') for model_name in models: print(f'>>>> 初始化模型 {model_name}') start_time = time.time() init_model(model_name) print(f'>>>> 初始化模型 {model_name} 完成 ✅，耗时 {time.time() - start_time:.4f} 秒') print(f'>>>> 初始化所有模型完成 ✅') return '初始化成功！' except Exception as e: logger.error(f'Data init failed: {str(e)}') return f'初始化失败: {str(e)}' with gr.Blocks(title='Neko Sticker Search 🔍', css='.gradio-container img { width: 200px !important; height: 200px !important; object-fit: contain; }') as demo: with gr.Row(): search_input = gr.Textbox(label='搜索关键词') search_button = gr.Button('搜索') headers = ['序号'] + [f'🧊{model.split("/")[-1]}' for i, model in enumerate(models)] results_table = gr.Dataframe( headers=headers, datatype=['number'] + ['markdown'] * len(models), row_count=5, col_count=len(models) + 1 ) status_box = gr.Textbox(label='状态', interactive=False) refresh_button = gr.Button('刷新数据') refresh_button.click(fn=init_collections, outputs=status_box) # 由于这里只是简单的搜索操作，可以直接使用同步方式调用 search_button.click( fn=search_all_models, inputs=[search_input], outputs=results_table ) return demo if __name__ == '__main__': # 提前加载所有模型 start_time = time.time() for index, model_name in enumerate(models): try: start_time = time.time() searchers[model_name] = StickerSearcher(SentenceTransformerModel(model_name)) print(f'>>>> 预加载模型 {model_name} 完成 ✅, 耗时 {time.time() - start_time:.4f} 秒') except Exception as e: logger.error(f'Model preload failed: {model_name} | Error: {str(e)}') logger.info(f'>>>> 预加载模型完成 ✅: {models}, 耗时 {time.time() - start_time:.4f} 秒') demo = create_gradio_ui() demo.launch()