LangChain Custom Embedding Model: Subclass Embeddings

How do I create a custom embedding class in LangChain?

LangChain vector stores only need an object with two methods. The abstract base class Embeddings in langchain_core.embeddings defines them: embed_documents(texts) returns one vector per input text, and embed_query(text) returns a single vector for a search string.

The example below wraps a local sentence-transformers model, such as a fine-tuned checkpoint saved on disk. The model loads once in the constructor, and encode(...).tolist() converts numpy output to plain lists of Python floats.

from typing import List
from langchain_core.embeddings import Embeddings
from sentence_transformers import SentenceTransformer


class MyEmbeddings(Embeddings):
    def __init__(self, model_name_or_path: str = "./my-finetuned-model",
                 device: str = "cpu", batch_size: int = 32):
        # Load ONCE. Loading per call is the most common performance bug.
        self.model = SentenceTransformer(model_name_or_path, device=device)
        self.batch_size = batch_size

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        vecs = self.model.encode(
            texts,
            batch_size=self.batch_size,
            normalize_embeddings=True,
            show_progress_bar=False,
        )
        return vecs.tolist()

    def embed_query(self, text: str) -> List[float]:
        vec = self.model.encode(text, normalize_embeddings=True)
        return vec.tolist()

How do I plug the class into Chroma?

Create one instance and hand it to the vector store. The parameter name depends on which entry point you use, and this trips up many people.

from_documents and from_texts use embedding. The Chroma constructor uses embedding_function. In both cases you pass the instance, never the class and never a bare function.

1. Install the packages: pip install langchain-core langchain-chroma sentence-transformers.
2. Build your documents as a list of Document objects.
3. Create the store with from_documents (embeds and persists) or reopen it later with the constructor.
4. Query with similarity_search or similarity_search_with_score.

from langchain_core.documents import Document
from langchain_chroma import Chroma

emb = MyEmbeddings("./my-finetuned-model")

docs = [
    Document(page_content="Chroma stores vectors on disk.", metadata={"src": "a"}),
    Document(page_content="Bananas are rich in potassium.", metadata={"src": "b"}),
]

# Build: keyword is `embedding`
db = Chroma.from_documents(
    docs,
    embedding=emb,
    persist_directory="./chroma_db",
    collection_name="kb",
    collection_metadata={"hnsw:space": "cosine"},
)

# Reopen later: keyword is `embedding_function`
db = Chroma(
    collection_name="kb",
    embedding_function=emb,
    persist_directory="./chroma_db",
)

print(db.similarity_search("where are vectors saved?", k=1))

Which import paths apply to which version?

LangChain split into packages, so older tutorials import from places that are now deprecated or removed. Use the current paths for new code and know the old ones when reading legacy examples.

The collection_metadata hnsw:space key is read by Chroma's collection creation. Recent Chroma releases also accept a configuration argument for the index space, so if your installed version warns about hnsw:space, follow its message.

Why do I get TypeError or dimension errors?

Most failures come from a short list of causes. Python enforces abstract methods at instantiation, and Chroma fixes a collection's vector size on the first insert, so errors surface early and are usually literal.

Use this table to match the message to the fix.

How do I wrap an ONNX model or a local HTTP server?

The pattern is identical: the two methods are the only contract. What changes is how you produce a vector.

For an HTTP backend such as Ollama, TEI (text-embeddings-inference) or llama.cpp's server, call the endpoint inside embed_documents. Send the whole batch in one request when the server supports it, and reuse a requests.Session to avoid reconnecting every call. The Ollama /api/embed route accepts a list under input and returns embeddings; verify the route and payload of your server version before relying on it.

For a raw transformers or ONNX model, you must do pooling yourself because the model returns token-level states. Use attention-mask-aware mean pooling, then L2-normalize.

import requests
import numpy as np
from typing import List
from langchain_core.embeddings import Embeddings


class OllamaHTTPEmbeddings(Embeddings):
    def __init__(self, model="nomic-embed-text",
                 url="http://localhost:11434/api/embed", timeout=60):
        self.model, self.url, self.timeout = model, url, timeout
        self.session = requests.Session()

    def _post(self, inputs: List[str]) -> List[List[float]]:
        r = self.session.post(self.url, json={"model": self.model, "input": inputs},
                              timeout=self.timeout)
        r.raise_for_status()
        return r.json()["embeddings"]

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        out = []
        for i in range(0, len(texts), 64):   # keep request size bounded
            out.extend(self._post(texts[i:i + 64]))
        return out

    def embed_query(self, text: str) -> List[float]:
        return self._post([text])[0]


def mean_pool_normalize(last_hidden, attention_mask):
    # last_hidden: (batch, seq, dim) ; attention_mask: (batch, seq)
    mask = attention_mask[..., None].astype(last_hidden.dtype)
    summed = (last_hidden * mask).sum(axis=1)
    counts = np.clip(mask.sum(axis=1), 1e-9, None)
    pooled = summed / counts
    pooled /= np.linalg.norm(pooled, axis=1, keepdims=True)
    return pooled   # call .tolist() before returning from embed_documents

What about query prefixes, normalization, batching and async?

Models such as E5 and BGE are trained with asymmetric inputs. E5 expects "query: " before search strings and "passage: " before documents. BGE English v1.5 models use an instruction prefix for queries ("Represent this sentence for searching relevant passages: ") and none for documents. This is exactly why embed_query and embed_documents are separate methods: they may legitimately differ.

Normalization matters for distance. Chroma defaults to squared L2. With unit-length vectors, L2 and cosine give the same ranking, but set hnsw:space to cosine in collection_metadata so scores are interpretable. Without normalization, use cosine or inner product deliberately.

For speed, load once, raise batch_size (32 to 128 on GPU, depending on memory), and set device="cuda" when available. The base class provides default aembed_documents and aembed_query that run your sync methods in an executor; override them with a real async client (for example httpx.AsyncClient) when your backend is remote.

class E5Embeddings(MyEmbeddings):
    def embed_documents(self, texts):
        return super().embed_documents([f"passage: {t}" for t in texts])

    def embed_query(self, text):
        return super().embed_query(f"query: {text}")

How do I test it, and when do I not need a custom class?

Before indexing thousands of chunks, run a quick sanity check: the vector length must match the model's documented dimension, and a related pair must score higher than an unrelated one.

If your model is a standard Hugging Face checkpoint, a built-in wrapper may be enough: HuggingFaceEmbeddings from langchain_huggingface (sentence-transformers under the hood, with model_kwargs and encode_kwargs such as normalize_embeddings), OllamaEmbeddings from langchain_ollama, or FastEmbedEmbeddings from langchain_community. Write a custom class when you need prefixes, custom pooling, ONNX, authentication or a private service.

import numpy as np

emb = MyEmbeddings("./my-finetuned-model")
q = emb.embed_query("how do I reset my password?")
d = emb.embed_documents(["Steps to reset a password", "Banana bread recipe"])

assert isinstance(q, list) and isinstance(q[0], float)
print("dim:", len(q))
sims = np.array(d) @ np.array(q)   # cosine if normalized
print(sims)
assert sims[0] > sims[1]

# Record the model so a later change is caught
import json
json.dump({"model": "./my-finetuned-model", "dim": len(q)},
          open("./chroma_db/embedding_meta.json", "w"))