Skip to content

Model2Vec Embedding Engine

Source: R/sentence_embeddings-model2vec.R
embedding_model2vec.Rd

embedding_model2vec() creates an engine for step_sentence_embeddings() that embeds text with a Model2Vec model. The model is a table of token vectors that is read from disk and used in R, so it needs neither Python nor a server, and text never leaves your machine.

Usage

embedding_model2vec(path, normalize = TRUE, max_tokens = 512)

Arguments

path

A single string giving the path to a folder holding a Model2Vec model, which must contain the files model.safetensors and vocab.txt. download_embeddings_model2vec() downloads such a folder. The path is expanded and stored in the engine.

normalize

A single logical. If TRUE, each embedding is scaled to have unit length, as the Model2Vec Python package does by default. Defaults to TRUE.

max_tokens

A single positive whole number giving the maximum number of tokens of a text that are used. Later tokens are ignored. Defaults to 512, as in the Model2Vec Python package.

Value

An object of class sentence_embedding_engine.

Details

The embedding of a text is the mean of the vectors of its tokens, scaled to unit length if normalize = TRUE. This is what Model2Vec does. The text is tokenized with the WordPiece vocabulary of the model using wordpiece::wordpiece_tokenize(), tokens that are not in the vocabulary ("[UNK]") are ignored, and no special tokens are added. A text without any known tokens gets a vector of zeros.

This follows the way Model2Vec builds an embedding, but it is not guaranteed to match the Python model2vec package exactly. Tokenizer details such as the handling of accents or unusual characters may differ, so compare a few texts with the Python package if you need identical values. Only models that use a WordPiece vocabulary, such as the English potion models, are supported. Models with a different tokenizer, with quantized weights, or with extra weights for the tokens can't be used.

The engine stores the path to the model folder, not the model. A prepped recipe therefore needs the model files to be at the same path where it is baked. The model is read the first time it is used and is then kept for the rest of the R session. Each parallel worker reads its own copy.

Requires the wordpiece, safetensors and torch packages. The torch package is only used to read the file, and the model is not run with it.

See also

step_sentence_embeddings(), download_embeddings_model2vec(), and vignette("static-embeddings", package = "textrecipes") for using the same model with step_word_embeddings().

Other Embedding Engines: embedding_custom(), embedding_huggingface(), embedding_ollama()

Examples

if (FALSE) {
# Downloads the model files
library(recipes)

url <- paste0(
  "https://huggingface.co/minishlab/potion-base-8M/resolve/",
  "bf8b056651a2c21b8d2565580b8569da283cab23"
)
model_dir <- download_embeddings_model2vec(url)

sample_data <- tibble(
  text = c("The cat sat.", "A dog ran.", "Stocks fell sharply today.")
)

rec <- recipe(~text, data = sample_data) |>
  step_sentence_embeddings(text, engine = embedding_model2vec(model_dir))

rec |>
  prep() |>
  bake(new_data = NULL)
}