weightgain

Fine-tune any embedding model in under a minute. Even models from OpenAI, Cohere, Voyage, etc.

It works by training an adapter that sits on top of the model, instead of modifying the model itself. This adapter transforms your embeddings after they're generated to boost retrieval accuracy and overall RAG performance.

Weightgain lets you train an adapter in just a couple lines of code, even if you don't have a dataset.

Installation

> pip install weightgain

Quickstart

from weightgain import Dataset, Adapter

# Generate a dataset (or supply your own)
dataset = Dataset.from_synthetic_chunks(
    prompt="Chunks of code from an arbitrary Python codebase.",
    llm="openai/gpt-4o-mini",
)

# Train the adapter
adapter = Adapter("openai/text-embedding-3-large")
adapter.fit(dataset)

# Apply the adapter
new_embeddings = adapter.transform(old_embeddings)

Usage

Choosing an Embedding Model

Weightgain wraps LiteLLM to provide access to models. You can fine-tune models from OpenAI, Cohere, Voyage, and more. Here's the full list of supported models.

Building the Dataset

You need a dataset of [query, chunk] pairs to get started. A chunk is a retrieval result, e.g. a code snippet or excerpt from a document. You can either generate a synthetic dataset or supply your own.

If you already have chunks:

from weightgain import Dataset

chunks = [...] # list of strings
dataset = Dataset.from_chunks(
    chunks,
    llm="openai/gpt-4o-mini",
    n_queries_per_chunk=1
)

This will use gpt-4o-mini (or whatever LiteLLM model you want) to generate 1 query per chunk.

If you don't have chunks:

dataset = Dataset.from_synthetic_chunks(
    prompt="Chunks of code from an arbitrary Python codebase.",
    llm="openai/gpt-4o-mini",
    n_chunks=25,
    n_queries_per_chunk=1
)

This will generate chunks using the prompt, and then generate 1 query per chunk.

If you have queries and chunks:

qa_pairs = [...] # list of (str, str) tuples
dataset = Dataset.from_pairs(qa_pairs, model)

Training the Adapter

from weightgain import Adapter

adapter = Adapter.fit(
    dataset,
    batch_size=25,
    max_epochs=50,
    learning_rate=100.0,
    dropout=0.0
)

After training, you can generate a report with various plots (training loss, cosine similarity distributions before/after training, etc.):

adapter.show_report()

Using the Adapter

old_embeddings = [...] # list of vectors
new_embeddings = adapter.transform(old_embeddings)

Behind the scenes, an adapter is just a matrix of weights that you can multiply your embeddings with. You can access this matrix like so:

adapter.matrix # returns numpy.ndarray

Roadmap

1. Add option to train an MLP instead of a linear layer
2. Add a method for easy hyperparameter search