如何使用 "套娃词向量" 模型提升RAG检索速度和效果(附论文+模型权重)

{"a": 1,"b": 2,"c": 3,..."z": 26,"the": 27," ": 28}

!pip install torch sentence_transformers einops

import torchfrom sentence_transformers import SentenceTransformer


device = "cuda" if torch.cuda.is_available() else"cpu"model = SentenceTransformer("nomic-ai/nomic-embed-text-v1.5",device=device,trust_remote_code=True,prompts={"search_query": "search_query: ","search_document": "search_document: ","classification": "classification: ","clustering": "clustering: ",},)





def embed_sentences(model: SentenceTransformer,sentences: list[str],prompt_name: str,matryoshka_dim: int,device: str,):assert matryoshka_dim <= 768, "maximum dimension for nomic-embed-text-v1.5 is 768"embeddings = model.encode(sentences, prompt_name=prompt_name, device=device, convert_to_tensor=True)embeddings = torch.nn.functional.layer_norm(embeddings, normalized_shape=(embeddings.shape[1],))embeddings = embeddings[:, :matryoshka_dim]embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)return embeddings.cpu()

matryoshka_dim = 64


wikipedia_texts = ["The dog (Canis familiaris or Canis lupus familiaris) is a domesticated descendant of the wolf.","Albert Einstein was born in Ulm in the Kingdom of Württemberg in the German Empire, on 14 March 1879.","Einstein excelled at physics and mathematics from an early age, and soon acquired the mathematical expertise normally only found in a child several years his senior.","Werner Karl Heisenberg was a German theoretical physicist, one of the main pioneers of the theory of quantum mechanics, and a principal scientist in the Nazi nuclear weapons program during World War II.","Steven Paul Jobs (February 24, 1955 - October 5, 2011) was an American businessman, inventor, and investor best known for co-founding the technology giant Apple Inc.","The cat (Felis catus), commonly referred to as the domestic cat or house cat, is the only domesticated species in the family Felidae.",]


question = ["Where was Albert Einstein born?"]


question_embedding = embed_sentences(model,sentences=question,prompt_name="search_query",matryoshka_dim=matryoshka_dim,device=device,)





document_embeddings = embed_sentences(model,sentences=wikipedia_texts,prompt_name="search_document",matryoshka_dim=matryoshka_dim,device=device,)print(f"document_embeddings.shape: {document_embeddings.shape}")print(f"question_embedding.shape:{question_embedding.shape}")

!pip install faiss-cpu

import faiss


index = faiss.IndexFlatIP(matryoshka_dim)index.add(document_embeddings)

print(index.ntotal)>> 6

distances, indices = index.search(question_embedding, k=6)print(indices)print(distances)>> [[1 2 3 4 0 5]]>> [[0.96335280.729192 0.63353264 0.62068397 0.512541 0.43155164]]

# d=64 的结果print(question)print(wikipedia_texts[1])>> ['阿尔伯特·爱因斯坦出生在哪里？']>> '阿尔伯特·爱因斯坦出生在德意志帝国的维尔滕贝格王国乌尔姆，于1879年3月14日。'

# 结果与 d=768print(indices)print(distances)>> [[1 2 4 3 0 5]]>> [[0.92466116 0.645744 0.54405797 0.54004824 0.39331824 0.37972206]]