-quality embeddings for a wide range of NLP tasks.from typing import Annotated, List

import numpy as np

from sentence_transformers import SentenceTransformer

from structures import Document

from zenml import ArtifactConfig, log_artifact_metadata, step

@step

def generate_embeddings(

split_documents: List[Document],

) -> Annotated[

List[Document], ArtifactConfig(name="documents_with_embeddings")

]:

try:

model = SentenceTransformer("sentence-transformers/all-MiniLM-L12-v2")

log_artifact_metadata(

artifact_name="embeddings",

metadata={

"embedding_type": "sentence-transformers/all-MiniLM-L12-v2",

"embedding_dimensionality": 384,

},

document_texts = [doc.page_content for doc in split_documents]

embeddings = model.encode(document_texts)

for doc, embedding in zip(split_documents, embeddings):

doc.embedding = embedding

return split_documents

except Exception as e:

logger.error(f"Error in generate_embeddings: {e}")

raise

We update the Document Pydantic model to include an embedding attribute that stores the embedding generated for each document. This allows us to associate the embeddings with the corresponding documents and use them for retrieval purposes in the RAG pipeline.

There are smaller embeddings models if we cared a lot about speed, and larger ones (with more dimensions) if we wanted to boost our ability to retrieve more relevant chunks. The model we're using here is on the smaller side, but it should work well for our use case. The embeddings generated by this model have a dimensionality of 384, which means that each embedding is represented as a 384-dimensional vector in the high-dimensional space.