Как найти ближайшие векторы встраивания?

Question

У меня есть 100 000 известных вложений, т.е.

      [emb_1, emb_2, ..., emb_100000]

Каждое из этих вложений получено из вложения предложений GPT-3 с размерностью 2048.

Моей задаче дано встраивание () найти ближайшие 10 вложений из приведенных выше100kвложение.

То, как я подхожу к этой проблеме, - это грубая сила.

Каждый раз, когда запрос просит найти ближайшие вложения, я сравниваюembedding_newс[emb_1, emb_2, ..., emb_100000]и получить оценку сходства.

Затем я делаю быструю сортировку оценки сходства, чтобы получить верхнюю10ближайшее вложение.

В качестве альтернативы я также думал об использовании Faiss.

Есть ли лучший способ добиться этого?

Answer 1

Я нашел решение с помощью Vector Database Lite (VDBLITE)

VDBLITE здесь: https://pypi.org/project/vdblite/

      import vdblite
from time import time
from uuid import uuid4
import sys
from pprint import pprint as pp


if __name__ == '__main__':
    vdb = vdblite.Vdb()
    dimension = 12    # dimensions of each vector                         
    n = 200    # number of vectors                   
    np.random.seed(1)             
    db_vectors = np.random.random((n, dimension)).astype('float32')
    print(db_vectors[0])
    for vector in db_vectors:
        info = {'vector': vector, 'time': time(), 'uuid': str(uuid4())}
        vdb.add(info)
    vdb.details()
    results = vdb.search(db_vectors[10])
    pp(results)

Похоже, он использует FAISS за кулисами.

Answer 2

Используя свою собственную идею, просто убедитесь, что вложения имеют матричную форму, вы можете легко использоватьnumpyдля этого. Это вычисляется за линейное время (в количестве вложений) и должно быть быстрым.

       import numpy as np
k = 10 # k best embeddings
emb_mat = np.stack([emb_1, emb_2, ..., emb_100000])
scores = np.dot(emb_mat, embedding_new)
best_k_ind = np.argpartition(scores, k)[-k:] 
top_k_emb = emb_mat[best_k_ind]   

10 лучших вложений будут найдены вtop_k_emb. Для общего решения внутри программного проекта вы можете рассмотреть Faiss от Facebook Research. Пример использования Faiss:

       d = 2048  # dimensionality of your embedding data
k = 10  # number of nearest neighbors to return
index = faiss.IndexFlatIP(d)
emb_list = [emb_1, emb_2, ..., emb_100000]
index.add(emb_list)
D, I = index.search(embedding_new, k)

Вы можете использоватьIndexFlatIPдля подобия внутреннего продукта, илиindexFlatL2для евклидова\L2-нормового расстояния. Чтобы обойти проблемы с памятью (данные>1M), обратитесь к этой замечательной инфографике Faiss шпаргалка на слайде num. 7