Как постепенно добавлять индекс в FAISS Python

Я использую Faiss для индексации вложений моих огромных наборов данных, встраивания сгенерированных из модели Берта. Я хочу добавлять вложения постепенно, он работает нормально, если я добавляю его только с помощью faiss.IndexFlatL2, но проблема в том, что при сохранении его размер слишком велик. Итак, я попробовал с faiss.IndexIVFPQ, но ему нужно обучить вложения, прежде чем я добавлю данные, поэтому я не могу добавлять их постепенно, мне нужно сначала вычислить все вложения, а затем обучить и добавить их, это проблема, потому что все данные должны храниться в ОЗУ, пока я не напишу его. Есть ли способ сделать это постепенно. Вот мой код:

          # It is working fine when using with IndexFlatL2
    def __init__(self, sentences, model):
        self.sentences = sentences
        self.model = model
        self.index = faiss.IndexFlatL2(768)

    def process_sentences(self):
        result = self.model(self.sentences)
        self.sentence_ids = []
        self.token_ids = []
        self.all_tokens = []
        for i, (toks, embs) in enumerate(tqdm(result)):
            # initialize all_embeddings for every new sentence (INCREMENTALLY)
            all_embeddings = []
            for j, (tok, emb) in enumerate(zip(toks, embs)):
                self.sentence_ids.append(i)
                self.token_ids.append(j)
                self.all_tokens.append(tok)
                all_embeddings.append(emb)

            all_embeddings = np.stack(all_embeddings) # Add embeddings after every sentence
            self.index.add(all_embeddings)

        faiss.write_index(self.index, "faiss_Model")

И при использовании с IndexIVFPQ:

         def __init__(self, sentences, model):
       self.sentences = sentences
       self.model = model
       self.quantizer = faiss.IndexFlatL2(768)
       self.index = faiss.IndexIVFPQ(self.quantizer, 768, 1000, 16, 8)

   def process_sentences(self):
       result = self.model(self.sentences)
       self.sentence_ids = []
       self.token_ids = []
       self.all_tokens = []
       all_embeddings = []
       for i, (toks, embs) in enumerate(tqdm(result)):
           for j, (tok, emb) in enumerate(zip(toks, embs)):
               self.sentence_ids.append(i)
               self.token_ids.append(j)
               self.all_tokens.append(tok)
               all_embeddings.append(emb)

       all_embeddings = np.stack(all_embeddings)
       self.index.train(all_embeddings) # Train
       self.index.add(all_embeddings) # Add to index
       faiss.write_index(self.index, "faiss_Model_mini")

0 ответов

Другие вопросы по тегам