Class SearchEngine

search_engine/search_engine.py:10–332

class SearchEngine:
    def __init__(self, embed_model_name='GVE'): # Alibaba-NLP/GVE-3B Alibaba-NLP/GVE-7B

        self.embed_model_name = embed_model_name

        self.embed_model_name_file = embed_model_name.replace('/', '_').replace('-','_')

        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        if 'GVE' in self.embed_model_name:
            from models.GVE import AutoModelForSentenceEmbeddingTriplet
            from models.GVE import VLProcessor
            if '3B' in self.embed_model_name:
                model_name_or_path = "Alibaba-NLP/GVE-3B"
                self.dimension = 2048
            elif '7B' in self.embed_model_name:
                model_name_or_path = "Alibaba-NLP/GVE-7B"
                self.dimension = 3584

            self.model = AutoModelForSentenceEmbeddingTriplet(model_name_or_path)
            self.processor = VLProcessor(model_name_or_path)
            
        elif 'bge-m3' in self.embed_model_name:
            from sentence_transformers import SentenceTransformer
            self.model = SentenceTransformer(self.embed_model_name)
            self.dimension = 1024
            
        elif 'Qwen3-VL' in self.embed_model_name or 'Qwen3_VL' in self.embed_model_name:
            from models.Qwen3_VL_Embedding.qwen3_vl_embedding import Qwen3VLEmbedder
            if '2B' in self.embed_model_name:
                self.dimension = 2048
            elif '8B' in self.embed_model_name:
                self.dimension = 4096
            self.model = Qwen3VLEmbedder(self.embed_model_name)
            

    def load_index(self, input_dir):
        # meta data
        with open(os.path.join(input_dir, f"{self.embed_model_name_file}_meta_data.json"), "r") as f:
            self.meta_data = json.load(f)
        assert self.meta_data["model_name"] == self.embed_model_name
        print(f"loading model {self.embed_model_name}; loading index from {input_dir}; loading index length {self.meta_data['num_vectors']}...")
        
        self.index = faiss.read_index(os.path.join(input_dir, f"{self.embed_model_name_file}_faiss_index.bin"))
        with open(os.path.join(input_dir, f"{self.embed_model_name_file}_file_data_list.jsonl"), "r") as f:
            self.file_data_list = [json.loads(line) for line in f]
        assert self.meta_data["num_vectors"] == self.index.ntotal \
            and len(self.file_data_list) == self.index.ntotal, \
            f"meta_data_num_vectors: {self.meta_data['num_vectors']} and index_ntotal: {self.index.ntotal} mismatch \n or file_data_list_length: {len(self.file_data_list)} and index_ntotal: {self.index.ntotal} mismatch"

    def load_multi_index_corpus_together(self, input_dir_list):
        for input_dir in input_dir_list:
            if not hasattr(self, "meta_data"):
                with open(os.path.join(input_dir, f"{self.embed_model_name_file}_meta_data.json"), "r") as f:
                    self.meta_data = json.load(f)
                    num_vectors = self.meta_data['num_vectors']
                    print(f"Loading index from {input_dir}: {num_vectors}")