embetter

To accelerate embedding calculations (especially on large datasets), you can use Modal to run SentenceEncoder or ClipEncoder on remote hardware with a GPU. These encoders automatically detect and utilize available GPUs.

When using Modal, ensure your container image includes embetter[text] and the necessary dependencies.

import modal

stub = modal.Stub("example-get-started")
image = (modal.Image.debian_slim()
         .pip_install("simsity", "embetter[text]", "h5py")
         .run_commands("python -c 'from embetter.text import SentenceEncoder; SentenceEncoder()'"))

@stub.function(image=image, gpu="any")
def create(data):
    from embetter.text import SentenceEncoder
    return SentenceEncoder().transform(data)

@stub.local_entrypoint()
def main():
    data = read_text() # User defined
    X = create.call(data)
    print(f"Embedded shape: {X.shape}")

For applications where speed is critical and heavy pretrained models are overkill, you can use learn_lite_doc_embeddings. This technique uses a combination of TfidfVectorizer and TruncatedSVD to create fast, lightweight representations.

Key features:

• Extremely fast training and transformation.
• Can be saved to disk during the training call using the path parameter.
• Useful for quick visualization (e.g., with UMAP and plot_text).

import srsly
from umap import UMAP
from cluestar import plot_text
from embetter.text import learn_lite_doc_embeddings

# Load texts
texts = [ex['text'] for ex in srsly.read_jsonl("datasets/new-dataset.jsonl")]

# Train embeddings and optionally save to disk
enc = learn_lite_doc_embeddings(texts, dim=300, path="stored/on/disk.emb")

# Transform and visualize
X_orig = enc.transform(texts)
X = UMAP().fit_transform(X_orig)
plot_text(X, texts)

To avoid expensive re-calculations of embeddings (especially when using external providers), use the embetter.utils.cached helper. This integrates with diskcache to store results on disk.

When using cached, provide a unique string name for the cache. Subsequent calls to transform with the same input will be significantly faster as they fetch from the cache instead of re-encoding. You can also access precalculated embeddings directly using a diskcache.Cache instance with the same name used in cached.

from embetter.text import SentenceEncoder
from embetter.utils import cached

# Wrap the encoder with a cache name
encoder = cached("sentence-enc", SentenceEncoder('all-MiniLM-L6-v2'))

examples = [f"text {i}" for i in range(10_000)]

# First run: slow (writes to cache)
encoder.transform(examples)

# Second run: fast (reads from cache)
encoder.transform(examples)

# Direct access via diskcache
from diskcache import Cache
cache = Cache("sentence-enc")
embedding = cache["this is a pretty long text, which is more expensive 0"]

To create a pipeline that extracts text from a pandas DataFrame and converts it into embeddings using Sentence-Transformers, use ColumnGrabber followed by SentenceEncoder.

Note: This requires the sbert extra: pip install 'embetter[sbert]'.

import pandas as pd
from sklearn.pipeline import make_pipeline 
from sklearn.linear_model import LogisticRegression
from embetter.grab import ColumnGrabber
from embetter.text import SentenceEncoder

# Pipeline to transform text column into embeddings
text_emb_pipeline = make_pipeline(
  ColumnGrabber("text"),
  SentenceEncoder('all-MiniLM-L6-v2')
)

dataf = pd.DataFrame({
  "text": ["positive sentiment", "super negative"],
  "label_col": ["pos", "neg"]
})

# Get embeddings
X = text_emb_pipeline.fit_transform(dataf, dataf['label_col'])

# Full pipeline for classification
text_clf_pipeline = make_pipeline(
  ColumnGrabber("text"),
  SentenceEncoder('all-MiniLM-L6-v2'),
  LogisticRegression()
)
text_clf_pipeline.fit(dataf, dataf['label_col']).predict(dataf)