R/embed.R

In semanticfa: Semantic Factor Analysis of Language Model Embeddings

# Default sbert embedding model. Keep in sync with the `model =` defaults in
# sfa() and sfa_embed(); print.sfa() uses this to decide whether to show the
# "larger model" upgrade hint.
.SFA_DEFAULT_MODEL <- "Qwen/Qwen3-Embedding-0.6B"

#' @keywords internal
# Resolve a NULL model to the right default for the chosen backend, so that
# embed = "openai" does not inherit the sbert (Qwen) default.
.resolve_embed_model <- function(embed, model) {
  if (!is.null(model)) return(model)
  if (is.character(embed)) {
    switch(embed,
           sbert  = .SFA_DEFAULT_MODEL,
           openai = "text-embedding-3-small",
           .SFA_DEFAULT_MODEL)
  } else {
    .SFA_DEFAULT_MODEL
  }
}

#' Embed Item Text with a Language Model
#'
#' Computes embeddings for a vector of item text using a sentence-transformer
#' or other embedding backend.
#'
#' @param embed Embedding backend: \code{"sbert"} (default, via
#'   \code{reticulate}), \code{"openai"} (via \code{httr2}), or a function
#'   taking a character vector and returning a numeric matrix.
#' @param model Model name passed to the backend. If \code{NULL} (default), a
#'   backend-appropriate default is used: \code{"Qwen/Qwen3-Embedding-0.6B"} for
#'   \code{"sbert"} and \code{"text-embedding-3-small"} for \code{"openai"}.
#'   Larger embedding models recover factor structure more accurately; see
#'   \code{\link{sfa}}.
#' @param cache Logical: cache embeddings in
#'   \code{tools::R_user_dir("semanticfa", "cache")}? Default \code{TRUE}.
#' @param ... Additional arguments passed to the embedding backend function.
#'
#' @returns A numeric matrix (n_items x embedding_dim). Rownames are the item
#'   codes when \code{items} is a data frame with a \code{code} column,
#'   otherwise the item text.
#'
#' @param items Character vector of item text, or a data frame with an
#'   \code{item}/\code{text} column (and optionally a \code{code} column, used
#'   as rownames so short codes flow through to plots such as
#'   \code{\link{sfa_corplot}}).
#' @export
sfa_embed <- function(items, embed = "sbert", model = NULL,
                      cache = TRUE, ...) {
  row_labels <- NULL
  if (is.data.frame(items)) {
    resolved   <- .resolve_items(items)
    row_labels <- resolved$codes        # the 'code' column, when present
    items      <- resolved$items        # the item text
  }
  if (is.null(row_labels)) row_labels <- items

  if (is.function(embed)) {
    emb <- embed(items, ...)
    if (!is.matrix(emb) || !is.numeric(emb)) {
      stop("Custom embed function must return a numeric matrix.", call. = FALSE)
    }
    if (nrow(emb) != length(items)) {
      stop("Custom embed function returned ", nrow(emb), " rows for ",
           length(items), " items.", call. = FALSE)
    }
    rownames(emb) <- row_labels
    return(emb)
  }

  embed <- match.arg(embed, c("sbert", "openai"))
  model <- .resolve_embed_model(embed, model)

  if (cache) {
    key <- .cache_key(items, model, embed)
    cache_dir <- tools::R_user_dir("semanticfa", "cache")
    cache_file <- file.path(cache_dir, paste0(key, ".rds"))
    if (file.exists(cache_file)) {
      cached <- readRDS(cache_file)
      if (is.matrix(cached) && nrow(cached) == length(items)) {
        rownames(cached) <- row_labels    # honor codes vs text for this call
        return(cached)
      }
    }
  }

  emb <- switch(embed,
    sbert  = .embed_sbert(items, model, ...),
    openai = .embed_openai(items, model, ...)
  )

  rownames(emb) <- row_labels

  if (cache) {
    if (!dir.exists(cache_dir)) dir.create(cache_dir, recursive = TRUE)
    saveRDS(emb, cache_file)
  }

  emb
}

#' Clear Embedding Cache
#'
#' Removes all cached embedding files created by [sfa_embed()].
#'
#' @returns Invisible \code{NULL}.
#' @export
sfa_clear_cache <- function() {
  cache_dir <- tools::R_user_dir("semanticfa", "cache")
  if (dir.exists(cache_dir)) {
    files <- list.files(cache_dir, full.names = TRUE)
    file.remove(files)
  }
  invisible(NULL)
}

#' @keywords internal
# SHA-256 over the items + model + backend (digest is in Imports): a cache hit
# must identify the exact request, so no weaker fallback hash is offered.
.cache_key <- function(items, model, backend = "sbert") {
  digest::digest(list(items = items, model = model, backend = backend),
                 algo = "sha256")
}

#' Provision the Python Environment for Embedding
#'
#' Declares and installs the Python packages needed by the \code{"sbert"}
#' embedding backend and the default \code{\link{sfa_nli_matrix}} classifier
#' (\code{sentence-transformers}, which pulls in \code{torch} and
#' \code{transformers}). With \pkg{reticulate} (>= 1.41) these requirements are
#' also declared automatically on first use via \code{reticulate::py_require()},
#' so calling this is optional --- it is handy for provisioning ahead of time
#' (e.g. on a machine with internet before running offline) or into a specific
#' environment.
#'
#' @param packages Character vector of Python packages to require/install.
#' @param ... Passed to \code{reticulate::py_install()} (e.g. \code{envname},
#'   \code{method}).
#'
#' @returns Invisible \code{NULL}.
#' @examples
#' \dontrun{
#' # one-time setup of the Python embedding environment
#' sfa_install_python()
#' }
#' @export
sfa_install_python <- function(packages = "sentence-transformers", ...) {
  .sfa_py_require(packages)
  reticulate::py_install(packages, ...)
  invisible(NULL)
}

#' @keywords internal
.sfa_py_require <- function(packages) {
  if ("py_require" %in% getNamespaceExports("reticulate")) {
    try(reticulate::py_require(packages), silent = TRUE)
  }
  invisible(NULL)
}

# One resident sentence-transformer at a time. Loading a 27B naming model
# next to an 8B extraction model would exceed a single 80 GB GPU, so loading
# a different model evicts the previous one and releases its GPU memory.
.sfa_encoder_env <- new.env(parent = emptyenv())

#' @keywords internal
.sfa_release_encoder <- function() {
  if (!is.null(.sfa_encoder_env$encoder)) {
    .sfa_encoder_env$encoder <- NULL
    .sfa_encoder_env$key <- NULL
    gc(verbose = FALSE)
    try({
      py_gc <- reticulate::import("gc")
      py_gc$collect()
      torch <- reticulate::import("torch")
      if (torch$cuda$is_available()) torch$cuda$empty_cache()
    }, silent = TRUE)
  }
  invisible(NULL)
}

#' @keywords internal
.embed_sbert <- function(items, model, ...) {
  .sfa_py_require("sentence-transformers")
  st <- tryCatch(
    reticulate::import("sentence_transformers"),
    error = function(e) {
      stop(
        "Python 'sentence-transformers' could not be loaded (",
        conditionMessage(e), ").\n",
        "Provision the Python environment with sfa_install_python(), or pass ",
        "precomputed embeddings: sfa(items, embeddings = your_matrix)",
        call. = FALSE
      )
    }
  )
  torch <- reticulate::import("torch")
  device <- if (torch$cuda$is_available()) {
    "cuda"
  } else if (torch$backends$mps$is_available()) {
    "mps"
  } else {
    "cpu"
  }
  # Optional weight dtype, e.g. options(semanticfa.torch_dtype = "bfloat16").
  # Large models (the 27B naming encoder) do not fit common GPUs at float32.
  dtype <- getOption("semanticfa.torch_dtype", NULL)
  if (!is.null(dtype)) {
    dtype <- match.arg(dtype, c("float16", "bfloat16", "float32"))
  }
  key <- paste(model, device, dtype %||% "default", sep = "|")
  if (!identical(.sfa_encoder_env$key, key)) {
    .sfa_release_encoder()
    encoder <- tryCatch(
      .sfa_load_st(st, torch, model, device, dtype),
      error = function(e) {
        # Some text-only checkpoints of multimodal families are misrouted by
        # sentence-transformers/transformers through a processor that demands
        # an image component. Fall back to a plain transformers pipeline that
        # reproduces the model's own modules.json (last-token pooling + L2
        # normalization); the pooling is attention-mask based, so it is
        # padding-side safe.
        message("sentence-transformers could not load '", model, "' (",
                conditionMessage(e), "); using the plain transformers ",
                "fallback (last-token pooling + L2 normalization).")
        try(reticulate::py_clear_last_error(), silent = TRUE)
        gc(verbose = FALSE)
        try({
          reticulate::import("gc")$collect()
          if (torch$cuda$is_available()) torch$cuda$empty_cache()
        }, silent = TRUE)
        .sfa_load_manual(model, device, dtype)
      }
    )
    .sfa_encoder_env$encoder <- encoder
    .sfa_encoder_env$key <- key
  }
  encoder <- .sfa_encoder_env$encoder
  emb <- if (inherits(encoder, "sfa_manual_encoder")) {
    .sfa_encode_manual(encoder, items)
  } else {
    encoder$encode(items, show_progress_bar = FALSE)
  }
  emb_r <- reticulate::py_to_r(emb)
  # a single item is returned as a 1-D array/vector; reshape to 1 x dim
  # (length(dim) is 0 for a vector and 1 for a 1-D array -- both != 2)
  if (length(dim(emb_r)) != 2L) {
    emb_r <- matrix(as.numeric(emb_r), nrow = length(items), byrow = TRUE)
  }
  storage.mode(emb_r) <- "double"
  emb_r
}

#' @keywords internal
.sfa_load_st <- function(st, torch, model, device, dtype) {
  if (is.null(dtype)) {
    return(st$SentenceTransformer(model, device = device))
  }
  torch_dtype <- switch(dtype,
    float16  = torch$float16,
    bfloat16 = torch$bfloat16,
    float32  = torch$float32
  )
  # transformers >= 5 renamed from_pretrained's torch_dtype to dtype; try
  # the current name first, then the older one.
  tryCatch(
    st$SentenceTransformer(model, device = device,
      model_kwargs = reticulate::dict(dtype = torch_dtype)),
    error = function(e) {
      st$SentenceTransformer(model, device = device,
        model_kwargs = reticulate::dict(torch_dtype = torch_dtype))
    }
  )
}

#' @keywords internal
.sfa_load_manual <- function(model, device, dtype) {
  tr <- reticulate::import("transformers")
  tok <- tr$AutoTokenizer$from_pretrained(model)
  # eager attention matches how the candidate pools were embedded (attention
  # backends differ slightly in bf16 numerics).
  mod <- if (identical(device, "cuda")) {
    # device_map streams shards straight to the GPU but needs accelerate;
    # without it, load on CPU and move (higher host-RAM peak, same result).
    tryCatch(
      tr$AutoModel$from_pretrained(model, dtype = dtype %||% "auto",
                                   attn_implementation = "eager",
                                   device_map = device),
      error = function(e)
        tr$AutoModel$from_pretrained(model, dtype = dtype %||% "auto",
                                     attn_implementation = "eager")$to(device)
    )
  } else {
    tr$AutoModel$from_pretrained(model, dtype = dtype %||% "auto",
                                 attn_implementation = "eager")$to(device)
  }
  mod <- mod$eval()
  structure(list(tokenizer = tok, model = mod, device = device),
            class = "sfa_manual_encoder")
}

# Python helper for the manual path: batched encode, last non-pad token per
# sequence, L2 normalization, float32 numpy output. Defined once per session.
#' @keywords internal
.sfa_manual_py <- function() {
  if (is.null(.sfa_encoder_env$manual_py)) {
    # All imports live INSIDE the function: py_run_string(local = TRUE)
    # executes with a separate locals dict, so module-level bindings are
    # invisible to the function body's global lookups.
    .sfa_encoder_env$manual_py <- reticulate::py_run_string("
def _sfa_manual_encode(model, tokenizer, texts, device, batch_size=8):
    import torch
    import numpy as np
    outs = []
    texts = list(texts)
    with torch.no_grad():
        for i in range(0, len(texts), batch_size):
            batch = tokenizer(texts[i:i + batch_size], padding=True,
                              truncation=True, max_length=512,
                              return_tensors='pt').to(device)
            h = model(**batch).last_hidden_state
            mask = batch['attention_mask']
            # Last REAL token per sequence, handling either padding side:
            # left-padded batches end on a real token at position -1; right-
            # padded batches need the mask-count index. (The mask-count
            # formula alone is WRONG under left padding, which decoder-family
            # tokenizers default to.)
            left_padded = bool((mask[:, -1].sum() == mask.shape[0]).item())
            if left_padded:
                pooled = h[:, -1]
            else:
                idx = mask.sum(dim=1) - 1
                pooled = h[torch.arange(h.size(0)), idx]
            # normalize in float32 (matches how the candidate pools were built)
            pooled = torch.nn.functional.normalize(pooled.float(), p=2, dim=1)
            outs.append(pooled.cpu().numpy())
    return np.vstack(outs)
", local = TRUE)
  }
  .sfa_encoder_env$manual_py
}

#' @keywords internal
.sfa_encode_manual <- function(encoder, items) {
  py <- .sfa_manual_py()
  py$`_sfa_manual_encode`(encoder$model, encoder$tokenizer,
                          as.list(items), encoder$device)
}

#' @keywords internal
.embed_openai <- function(items, model, ...) {
  if (!requireNamespace("httr2", quietly = TRUE)) {
    stop(
      "The 'openai' embedding backend requires the 'httr2' package.\n",
      "Install with: install.packages('httr2')",
      call. = FALSE
    )
  }
  api_key <- Sys.getenv("OPENAI_API_KEY", "")
  if (api_key == "") {
    stop("OPENAI_API_KEY environment variable is not set.", call. = FALSE)
  }

  resp <- httr2::request("https://api.openai.com/v1/embeddings") |>
    httr2::req_headers(
      Authorization = paste("Bearer", api_key),
      `Content-Type` = "application/json"
    ) |>
    httr2::req_body_json(list(input = as.list(items), model = model)) |>
    httr2::req_perform()

  body <- httr2::resp_body_json(resp)
  emb_list <- body$data
  emb_list <- emb_list[order(vapply(emb_list, `[[`, integer(1), "index"))]
  emb <- do.call(rbind, lapply(emb_list, function(x) unlist(x$embedding)))
  storage.mode(emb) <- "double"
  emb
}