View Source Rag.Embedding.Nx (rag v0.1.0)

Functions to generate embeddings using Nx.Serving.batched_run/2.

Summary

Functions

generate_embedding(generation, serving \\ Rag.EmbeddingServing)

Passes generation.query to serving to generate an embedding. Then, puts the embedding in generation.query_embedding.

generate_embedding(ingestion, serving \\ Rag.EmbeddingServing, text_key, embedding_key)

Passes the value of ingestion at text_key to serving to generate an embedding. Then, puts the embedding in ingestion at embedding_key.

generate_embeddings_batch(ingestions, serving \\ Rag.EmbeddingServing, text_key, embedding_key)

Passes all values of ingestions at text_key to serving to generate all embeddings in a single batch. Puts the embeddings in ingestions at embedding_key.

Functions

generate_embedding(generation, serving \\ Rag.EmbeddingServing)

@spec generate_embedding(Rag.Generation.t(), Nx.Serving.t()) :: Rag.Generation.t()

Passes generation.query to serving to generate an embedding. Then, puts the embedding in generation.query_embedding.

generate_embedding(ingestion, serving \\ Rag.EmbeddingServing, text_key, embedding_key)

@spec generate_embedding(
  map(),
  Nx.Serving.t(),
  text_key :: atom(),
  embedding_key :: atom()
) :: map()

Passes the value of ingestion at text_key to serving to generate an embedding. Then, puts the embedding in ingestion at embedding_key.

generate_embeddings_batch(ingestions, serving \\ Rag.EmbeddingServing, text_key, embedding_key)

@spec generate_embeddings_batch(
  [map()],
  Nx.Serving.t(),
  text_key :: atom(),
  embedding_key :: atom()
) :: [map()]

Passes all values of ingestions at text_key to serving to generate all embeddings in a single batch. Puts the embeddings in ingestions at embedding_key.