%%%------------------------------------------------------------------- %%% @doc ColBERT late interaction embedding provider %%% %%% Uses ColBERT models for multi-vector embeddings. Each document %%% produces multiple vectors (one per token) for fine-grained matching. %%% %%% Dependencies (transformers, torch) are installed automatically %%% in the managed venv on first use. %%% %%% == Configuration == %%% ``` %%% Config = #{ %%% model => "colbert-ir/colbertv2.0", %% Model name (default, 128 dims) %%% python => "python3", %% Python executable (default) %%% timeout => 120000 %% Timeout in ms (default) %%% }. %%% ''' %%% %%% == Multi-Vector Format == %%% Unlike single-vector embeddings, ColBERT produces a list of vectors: %%% ``` %%% [[0.1, 0.2, ...], [0.3, 0.4, ...], ...] %% One vector per token %%% ''' %%% %%% == Late Interaction == %%% ColBERT scoring uses MaxSim: %%% ``` %%% Score(Q, D) = sum(max(qi · dj for all dj in D) for all qi in Q) %%% ''' %%% This enables fine-grained token-level matching. %%% %%% == Supported Models == %%% - `"colbert-ir/colbertv2.0"' - Default, 128 dimensions %%% - `"answerdotai/answerai-colbert-small-v1"' - 96 dimensions, smaller %%% - `"jinaai/jina-colbert-v2"' - 128 dimensions, long context (8192 tokens) %%% %%% == Use Cases == %%% - Fine-grained semantic matching %%% - Passage retrieval with token-level scoring %%% - Question answering %%% %%% @end %%%------------------------------------------------------------------- -module(barrel_embed_colbert). -behaviour(barrel_embed_provider). %% Behaviour callbacks -export([ embed/2, embed_batch/2, dimension/1, name/0, init/1, available/1 ]). %% Multi-vector API -export([ embed_multi/2, embed_batch_multi/2, maxsim_score/2 ]). -define(DEFAULT_PYTHON, "python3"). -define(DEFAULT_MODEL, "colbert-ir/colbertv2.0"). -define(DEFAULT_TIMEOUT, 120000). -define(DEFAULT_DIMENSION, 128). %% Multi-vector type: list of token vectors -type multi_vector() :: [[float()]]. -export_type([multi_vector/0]). %%==================================================================== %% Behaviour Callbacks %%==================================================================== %% @doc Provider name. -spec name() -> atom(). name() -> colbert. %% @doc Get dimension for this provider. -spec dimension(map()) -> pos_integer(). dimension(Config) -> maps:get(dimension, Config, ?DEFAULT_DIMENSION). %% @doc Initialize the provider. -spec init(map()) -> {ok, map()} | {error, term()}. init(Config) -> Python = maps:get(python, Config, ?DEFAULT_PYTHON), Model = maps:get(model, Config, ?DEFAULT_MODEL), Timeout = maps:get(timeout, Config, ?DEFAULT_TIMEOUT), %% Use managed venv, auto-install deps Venv = get_managed_venv(colbert), %% Validate model (warning only) validate_model(Model), %% Build args for python -m barrel_embed Args = ["-m", "barrel_embed", "--provider", "colbert", "--model", Model], Opts = [ {timeout, Timeout}, {priv_dir, get_priv_dir()}, {venv, Venv} ], case barrel_embed_port_server:start_link(Python, Args, Opts) of {ok, Server} -> case barrel_embed_port_server:info(Server, Timeout) of {ok, #{dimensions := Dims}} -> {ok, Config#{ server => Server, dimension => Dims, timeout => Timeout }}; {ok, _} -> %% No dimensions in response, use default {ok, Config#{ server => Server, dimension => ?DEFAULT_DIMENSION, timeout => Timeout }}; {error, Reason} -> barrel_embed_port_server:stop(Server), {error, Reason} end; {error, Reason} -> {error, Reason} end. %% @doc Check if provider is available. -spec available(map()) -> boolean(). available(#{server := Server}) -> is_process_alive(Server); available(_Config) -> false. %% @doc Generate single-vector embedding (mean pooling of token vectors). %% Note: For ColBERT, use embed_multi/2 to get full multi-vector output. -spec embed(binary(), map()) -> {ok, [float()]} | {error, term()}. embed(Text, Config) -> case embed_multi(Text, Config) of {ok, MultiVec} -> {ok, mean_pool(MultiVec)}; {error, _} = Error -> Error end. %% @doc Generate single-vector embeddings for batch (mean pooling). -spec embed_batch([binary()], map()) -> {ok, [[float()]]} | {error, term()}. embed_batch(Texts, Config) -> case embed_batch_multi(Texts, Config) of {ok, MultiVecs} -> {ok, [mean_pool(MV) || MV <- MultiVecs]}; {error, _} = Error -> Error end. %%==================================================================== %% Multi-Vector API %%==================================================================== %% @doc Generate multi-vector embedding for a single text. %% Returns a list of token vectors. -spec embed_multi(binary(), map()) -> {ok, multi_vector()} | {error, term()}. embed_multi(Text, Config) -> case embed_batch_multi([Text], Config) of {ok, [MultiVec]} -> {ok, MultiVec}; {error, _} = Error -> Error end. %% @doc Generate multi-vector embeddings for multiple texts. -spec embed_batch_multi([binary()], map()) -> {ok, [multi_vector()]} | {error, term()}. embed_batch_multi(Texts, #{server := Server, timeout := Timeout}) -> barrel_embed_port_server:embed_multi_batch(Server, Texts, Timeout); embed_batch_multi(_Texts, _Config) -> {error, server_not_initialized}. %% @doc Calculate MaxSim score between query and document multi-vectors. %% This is the standard ColBERT scoring function. %% Score = sum(max(qi · dj for all dj in D) for all qi in Q) -spec maxsim_score(multi_vector(), multi_vector()) -> float(). maxsim_score(QueryVecs, DocVecs) -> lists:sum([max_dot_product(QVec, DocVecs) || QVec <- QueryVecs]). %%==================================================================== %% Internal Functions %%==================================================================== get_priv_dir() -> case code:priv_dir(barrel_embed) of {error, bad_name} -> "priv"; Dir -> Dir end. %% @private validate_model(Model) -> ModelBin = to_binary(Model), case is_known_model(ModelBin) of true -> ok; false -> error_logger:warning_msg( "Model ~s is not in the known list. " "It may still work if it's a valid ColBERT model.~n", [ModelBin] ) end. %% @private is_known_model(<<"colbert-ir/colbertv2.0">>) -> true; is_known_model(<<"answerdotai/answerai-colbert-small-v1">>) -> true; is_known_model(<<"jinaai/jina-colbert-v2">>) -> true; is_known_model(_) -> false. %% @private to_binary(S) when is_binary(S) -> S; to_binary(S) when is_list(S) -> list_to_binary(S). %% @private %% Mean pooling of token vectors to get single vector mean_pool([]) -> []; mean_pool(Vectors) -> N = length(Vectors), Dim = length(hd(Vectors)), %% Sum all vectors element-wise Sums = lists:foldl( fun(Vec, Acc) -> lists:zipwith(fun(A, B) -> A + B end, Vec, Acc) end, lists:duplicate(Dim, 0.0), Vectors ), %% Divide by N [S / N || S <- Sums]. %% @private %% Find maximum dot product between query vector and all doc vectors max_dot_product(QueryVec, DocVecs) -> DotProducts = [dot_product(QueryVec, DocVec) || DocVec <- DocVecs], lists:max(DotProducts). %% @private dot_product(V1, V2) -> lists:sum(lists:zipwith(fun(A, B) -> A * B end, V1, V2)). %% @private %% Get managed venv path and install deps for provider get_managed_venv(Provider) -> case application:get_env(barrel_embed, managed_venv_path) of {ok, Path} -> _ = barrel_embed_venv:install_deps(Provider), Path; undefined -> undefined end.