Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 0 additions & 18 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -91,24 +91,6 @@ doc2 = Nx.tensor([[0.5, 0.5]], type: :f32) # 1 token
scores = ExMaxsimCpu.maxsim_scores_variable(query, [doc1, doc2])
```

### Raw Binary API (Advanced)

For maximum performance when you already have raw binary data:

```elixir
# Query and docs as native-endian f32 binaries
query_bin = <<1.0::float-32-native, 0.0::float-32-native, ...>>
docs_bin = <<...>>

scores_bin = ExMaxsimCpu.maxsim_scores_raw(
query_bin, q_len, dim,
docs_bin, n_docs, d_len
)

# Parse scores
scores = for <<score::float-32-native <- scores_bin>>, do: score
```

## Performance Tuning

### Environment Variables
Expand Down
110 changes: 3 additions & 107 deletions lib/ex_maxsim_cpu.ex
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ defmodule ExMaxsimCpu do

## Usage with Nx Tensors

ExMaxsimCpu exposes an Nx-only public API.

# Query: [q_len, dim] tensor
query = Nx.tensor([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]], type: :f32)

Expand All @@ -20,15 +22,6 @@ defmodule ExMaxsimCpu do
scores = ExMaxsimCpu.maxsim_scores(query, docs)
# => #Nx.Tensor<f32[2]>

## Usage with Raw Binaries

For advanced use cases, you can use raw binaries directly:

query_bin = <<0.1::float-32-native, 0.2::float-32-native, ...>>
docs_bin = <<...>>

scores = ExMaxsimCpu.maxsim_scores_raw(query_bin, q_len, dim, docs_bin, n_docs, d_len)

## Performance Notes

- Uses Dirty CPU schedulers for compute-intensive operations
Expand Down Expand Up @@ -108,7 +101,7 @@ defmodule ExMaxsimCpu do
@spec maxsim_scores_variable(Nx.Tensor.t(), [Nx.Tensor.t()]) :: Nx.Tensor.t()
def maxsim_scores_variable(_query, []) do
raise ArgumentError,
"Empty document list not supported (Nx cannot create empty tensors). Use maxsim_scores_variable_raw/5 for empty lists."
"Empty document list not supported (Nx cannot create empty tensors)."
end

def maxsim_scores_variable(query, docs) when is_list(docs) do
Expand Down Expand Up @@ -145,103 +138,6 @@ defmodule ExMaxsimCpu do
Nx.from_binary(scores_bin, :f32)
end

@doc """
Compute MaxSim scores using raw binaries (advanced API).

This is a lower-level API for users who want to avoid Nx tensor overhead.

## Parameters

- `query_bin`: Binary containing query vectors as f32 values (native endian)
- `q_len`: Number of query tokens
- `dim`: Embedding dimension
- `docs_bin`: Binary containing document vectors as f32 values
- `n_docs`: Number of documents
- `d_len`: Number of tokens per document (must be uniform)

## Returns

Binary containing n_docs f32 scores.
"""
@spec maxsim_scores_raw(
binary(),
pos_integer(),
pos_integer(),
binary(),
pos_integer(),
pos_integer()
) ::
binary()
def maxsim_scores_raw(query_bin, q_len, dim, docs_bin, n_docs, d_len)
when is_binary(query_bin) and is_binary(docs_bin) and
is_integer(q_len) and q_len > 0 and
is_integer(dim) and dim > 0 and
is_integer(n_docs) and n_docs > 0 and
is_integer(d_len) and d_len > 0 do
expected_query_size = q_len * dim * 4
expected_docs_size = n_docs * d_len * dim * 4

if byte_size(query_bin) != expected_query_size do
raise ArgumentError,
"Query binary size mismatch: expected #{expected_query_size}, got #{byte_size(query_bin)}"
end

if byte_size(docs_bin) != expected_docs_size do
raise ArgumentError,
"Docs binary size mismatch: expected #{expected_docs_size}, got #{byte_size(docs_bin)}"
end

Nif.maxsim_scores_nif(query_bin, q_len, dim, docs_bin, n_docs, d_len)
end

@doc """
Compute MaxSim scores for variable-length documents using raw binaries (advanced API).

## Parameters

- `query_bin`: Binary containing query vectors as f32 values
- `q_len`: Number of query tokens
- `dim`: Embedding dimension
- `doc_bins`: List of binaries, each containing a document's vectors
- `doc_lens`: List of token counts for each document

## Returns

Binary containing n_docs f32 scores.
"""
@spec maxsim_scores_variable_raw(binary(), pos_integer(), pos_integer(), [binary()], [
pos_integer()
]) ::
binary()
def maxsim_scores_variable_raw(_query_bin, _q_len, _dim, [], []), do: <<>>

def maxsim_scores_variable_raw(query_bin, q_len, dim, doc_bins, doc_lens)
when is_binary(query_bin) and is_list(doc_bins) and is_list(doc_lens) do
expected_query_size = q_len * dim * 4

if byte_size(query_bin) != expected_query_size do
raise ArgumentError,
"Query binary size mismatch: expected #{expected_query_size}, got #{byte_size(query_bin)}"
end

if length(doc_bins) != length(doc_lens) do
raise ArgumentError, "doc_bins and doc_lens must have the same length"
end

# Validate each document binary size
Enum.zip(doc_bins, doc_lens)
|> Enum.with_index()
|> Enum.each(fn {{bin, len}, idx} ->
expected = len * dim * 4

if byte_size(bin) != expected do
raise ArgumentError,
"Doc #{idx} binary size mismatch: expected #{expected}, got #{byte_size(bin)}"
end
end)

Nif.maxsim_scores_variable_nif(query_bin, q_len, dim, doc_bins, doc_lens)
end

# Private helpers

Expand Down
2 changes: 1 addition & 1 deletion mix.exs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
defmodule ExMaxsimCpu.MixProject do
use Mix.Project

@version "0.2.0"
@version "0.3.0"
@source_url "https://github.com/mixedbread-ai/maxsim-cpu"

def project do
Expand Down
37 changes: 2 additions & 35 deletions test/ex_maxsim_cpu_test.exs
Original file line number Diff line number Diff line change
Expand Up @@ -92,41 +92,8 @@ defmodule ExMaxsimCpuTest do
test "handles empty list" do
query = Nx.tensor([[1.0, 0.0]], type: :f32)

# Nx doesn't support empty tensors, so we return an empty binary
scores =
ExMaxsimCpu.maxsim_scores_variable_raw(
Nx.to_binary(query),
1,
2,
[],
[]
)

assert scores == <<>>
end
end

describe "maxsim_scores_raw/6" do
test "works with raw binaries" do
# Query: 1 token, dim 2
query_bin = <<1.0::float-32-native, 0.0::float-32-native>>

# 1 doc: 1 token, dim 2
docs_bin = <<1.0::float-32-native, 0.0::float-32-native>>

scores_bin = ExMaxsimCpu.maxsim_scores_raw(query_bin, 1, 2, docs_bin, 1, 1)

assert byte_size(scores_bin) == 4
<<score::float-32-native>> = scores_bin
assert_in_delta(score, 1.0, 1.0e-5)
end

test "raises on size mismatch" do
query_bin = <<1.0::float-32-native>>
docs_bin = <<1.0::float-32-native>>

assert_raise ArgumentError, ~r/Query binary size mismatch/, fn ->
ExMaxsimCpu.maxsim_scores_raw(query_bin, 2, 2, docs_bin, 1, 1)
assert_raise ArgumentError, ~r/Empty document list not supported/, fn ->
ExMaxsimCpu.maxsim_scores_variable(query, [])
end
end
end
Expand Down
Loading