Use sqlite vector
This commit is contained in:
Claudio Ortolina
@@ -0,0 +1,22 @@
|
||||
defmodule SqliteVec.Ecto.Float32 do
|
||||
@moduledoc """
|
||||
`Ecto.Type` for `SqliteVec.Float32`
|
||||
"""
|
||||
use Ecto.Type
|
||||
|
||||
def type, do: :binary
|
||||
|
||||
def cast(value) do
|
||||
{:ok, SqliteVec.Float32.new(value)}
|
||||
end
|
||||
|
||||
def load(data) do
|
||||
{:ok, SqliteVec.Float32.from_binary(data)}
|
||||
end
|
||||
|
||||
def dump(%SqliteVec.Float32{} = vector) do
|
||||
{:ok, SqliteVec.Float32.to_binary(vector)}
|
||||
end
|
||||
|
||||
def dump(_), do: :error
|
||||
end
|
||||
@@ -0,0 +1,182 @@
|
||||
defmodule SqliteVec.Ecto.Query do
|
||||
@moduledoc """
|
||||
Macros for Ecto
|
||||
"""
|
||||
|
||||
@doc """
|
||||
Creates a bit vector
|
||||
"""
|
||||
defmacro vec_bit(vector) do
|
||||
quote do
|
||||
fragment("vec_bit(?)", type(^unquote(vector).data, :binary))
|
||||
end
|
||||
end
|
||||
|
||||
@doc """
|
||||
Creates an int8 vector
|
||||
"""
|
||||
defmacro vec_int8(vector) do
|
||||
quote do
|
||||
fragment("vec_int8(?)", type(^unquote(vector).data, :binary))
|
||||
end
|
||||
end
|
||||
|
||||
@doc """
|
||||
Creates a float32 vector
|
||||
"""
|
||||
defmacro vec_f32(vector) do
|
||||
quote do
|
||||
fragment("vec_f32(?)", type(^unquote(vector).data, :binary))
|
||||
end
|
||||
end
|
||||
|
||||
@doc """
|
||||
Calculates the L2 euclidian distance between vectors a and b. Only valid for float32 or int8 vectors.
|
||||
|
||||
Returns an error under the following conditions:
|
||||
- a or b are invalid vectors
|
||||
- a or b do not share the same vector element types (ex float32 or int8)
|
||||
- a or b are bit vectors. Use vec_distance_hamming() for distance calculations between two bitvectors.
|
||||
- a or b do not have the same length.
|
||||
"""
|
||||
# credo:disable-for-next-line Credo.Check.Readability.FunctionNames
|
||||
defmacro vec_distance_L2(a, b) do
|
||||
quote do
|
||||
fragment("vec_distance_L2(?, ?)", unquote(a), unquote(b))
|
||||
end
|
||||
end
|
||||
|
||||
@doc """
|
||||
Calculates the cosine distance between vectors a and b. Only valid for float32 or int8 vectors.
|
||||
|
||||
Returns an error under the following conditions:
|
||||
- a or b are invalid vectors
|
||||
- a or b do not share the same vector element types (ex float32 or int8)
|
||||
- a or b are bit vectors. Use vec_distance_hamming() for distance calculations between two bitvectors.
|
||||
- a or b do not have the same length
|
||||
"""
|
||||
defmacro vec_distance_cosine(a, b) do
|
||||
quote do
|
||||
fragment("vec_distance_cosine(?, ?)", unquote(a), unquote(b))
|
||||
end
|
||||
end
|
||||
|
||||
@doc """
|
||||
Calculates the hamming distance between two bitvectors a and b. Only valid for bitvectors.
|
||||
|
||||
Returns an error under the following conditions:
|
||||
- a or b are not bitvectors
|
||||
- a and b do not share the same length
|
||||
- Memory cannot be allocated
|
||||
"""
|
||||
defmacro vec_distance_hamming(a, b) do
|
||||
quote do
|
||||
fragment("vec_distance_hamming(?, ?)", unquote(a), unquote(b))
|
||||
end
|
||||
end
|
||||
|
||||
defmacro vec_match(a, b) do
|
||||
quote do
|
||||
fragment("? match ?", unquote(a), unquote(b))
|
||||
end
|
||||
end
|
||||
|
||||
@doc """
|
||||
Returns the number of elements in the given vector
|
||||
"""
|
||||
defmacro vec_length(vector) do
|
||||
quote do
|
||||
fragment("vec_length(?)", unquote(vector))
|
||||
end
|
||||
end
|
||||
|
||||
@doc """
|
||||
Returns the name of the type of `vector` as text
|
||||
"""
|
||||
defmacro vec_type(vector) do
|
||||
quote do
|
||||
fragment("vec_type(?)", unquote(vector))
|
||||
end
|
||||
end
|
||||
|
||||
@doc """
|
||||
Adds every element in vector a with vector b, returning a new vector c.
|
||||
Both vectors must be of the same type and same length.
|
||||
Only float32 and int8 vectors are supported.
|
||||
|
||||
An error is raised if either a or b are invalid, or if they are not the same type or same length.
|
||||
"""
|
||||
defmacro vec_add(a, b) do
|
||||
quote do
|
||||
fragment("vec_add(?, ?)", unquote(a), unquote(b))
|
||||
end
|
||||
end
|
||||
|
||||
@doc """
|
||||
Subtracts every element in vector a with vector b, returning a new vector c.
|
||||
Both vectors must be of the same type and same length.
|
||||
Only float32 and int8 vectors are supported.
|
||||
|
||||
An error is raised if either a or b are invalid, or if they are not the same type or same length.
|
||||
"""
|
||||
defmacro vec_sub(a, b) do
|
||||
quote do
|
||||
fragment("vec_sub(?, ?)", unquote(a), unquote(b))
|
||||
end
|
||||
end
|
||||
|
||||
@doc """
|
||||
Performs L2 normalization on the given vector.
|
||||
Only float32 vectors are currently supported.
|
||||
|
||||
Returns an error if the input is an invalid vector or not a float32 vector.
|
||||
"""
|
||||
defmacro vec_normalize(vector) do
|
||||
quote do
|
||||
fragment("vec_normalize(?)", unquote(vector))
|
||||
end
|
||||
end
|
||||
|
||||
@doc """
|
||||
Extract a subset of vector from the start element (inclusive) to the end element (exclusive).
|
||||
|
||||
This is especially useful for Matryoshka embeddings, also known as "adaptive length" embeddings.
|
||||
Use with vec_normalize() to get proper results.
|
||||
|
||||
Returns an error in the following conditions:
|
||||
- If vector is not a valid vector
|
||||
- If start is less than zero or greater than or equal to end
|
||||
- If end is greater than the length of vector, or less than or equal to start.
|
||||
- If vector is a bitvector, start and end must be divisible by 8.
|
||||
"""
|
||||
defmacro vec_slice(vector, start_index, end_index) do
|
||||
quote do
|
||||
fragment("vec_slice(?, ?, ?)", unquote(vector), unquote(start_index), unquote(end_index))
|
||||
end
|
||||
end
|
||||
|
||||
@doc """
|
||||
Represents a vector as JSON text.
|
||||
The input vector can be a vector BLOB or JSON text.
|
||||
|
||||
Returns an error if vector is an invalid vector, or when memory cannot be allocated.
|
||||
"""
|
||||
defmacro vec_to_json(vector) do
|
||||
quote do
|
||||
fragment("vec_to_json(?)", unquote(vector))
|
||||
end
|
||||
end
|
||||
|
||||
@doc """
|
||||
Quantize a float32 or int8 vector into a bitvector.
|
||||
For every element in the vector, a 1 is assigned to positive numbers and a 0 is assigned to negative numbers.
|
||||
These values are then packed into a bit vector.
|
||||
|
||||
Returns an error if vector is invalid, or if vector is not a float32 or int8 vector.
|
||||
"""
|
||||
defmacro vec_quantize_binary(vector) do
|
||||
quote do
|
||||
fragment("vec_quantize_binary(?)", unquote(vector))
|
||||
end
|
||||
end
|
||||
end
|
||||
@@ -0,0 +1,114 @@
|
||||
defmodule SqliteVec.Float32 do
|
||||
@moduledoc """
|
||||
A vector struct for float32 vectors.
|
||||
Vectors are stored as binaries in the endianness of the system.
|
||||
|
||||
> ### Consider endianness {: .warning}
|
||||
>
|
||||
> `SqliteVec.Float32.Vector` holds data in system endianness.
|
||||
> Therefore, the same vector data will be interpreted differently on another system with different endianness.
|
||||
> Moreover, you must consider endianness when converting the binary data directly to a list of numbers.
|
||||
|
||||
iex> v = SqliteVec.Float32.new([-1.0, 2.0])
|
||||
...> b = SqliteVec.Float32.to_binary(v)
|
||||
...> <<f1::float-32, f2::float-32>> = b
|
||||
...> [f1, f2]
|
||||
case System.endianness() do
|
||||
:big -> [-1.0, 2.0]
|
||||
:little -> [4.618539608568165e-41, 8.96831017167883e-44]
|
||||
end
|
||||
"""
|
||||
|
||||
@type t :: %__MODULE__{data: binary()}
|
||||
|
||||
defstruct [:data]
|
||||
|
||||
@doc """
|
||||
Creates a new vector from a vector, list, or tensor
|
||||
|
||||
The vector must be a `SqliteVec.Float32` vector.
|
||||
The list may contain any number but the values will be converted to f32 format.
|
||||
The tensor must have a rank of 1 and must be of type :f32.
|
||||
|
||||
## Examples
|
||||
iex> SqliteVec.Float32.new([1.0, 2.0])
|
||||
%SqliteVec.Float32{data: <<1.0::float-32-native, 2.0::float-32-native>>}
|
||||
|
||||
iex> v1 = SqliteVec.Float32.new([1, 2])
|
||||
...> SqliteVec.Float32.new(v1)
|
||||
%SqliteVec.Float32{data: <<1.0::float-32-native, 2.0::float-32-native>>}
|
||||
|
||||
iex> SqliteVec.Float32.new(Nx.tensor([1, 2], type: :f32))
|
||||
%SqliteVec.Float32{data: <<1.0::float-32-native, 2.0::float-32-native>>}
|
||||
"""
|
||||
def new(vector_or_list_or_tensor)
|
||||
|
||||
def new(%SqliteVec.Float32{} = vector) do
|
||||
vector
|
||||
end
|
||||
|
||||
def new(list) when is_list(list) do
|
||||
if list == [] do
|
||||
raise ArgumentError, "list must not be empty"
|
||||
end
|
||||
|
||||
bin = for v <- list, into: <<>>, do: <<v::float-32-native>>
|
||||
from_binary(<<bin::binary>>)
|
||||
end
|
||||
|
||||
if Code.ensure_loaded?(Nx) do
|
||||
def new(tensor) when is_struct(tensor, Nx.Tensor) do
|
||||
if Nx.rank(tensor) != 1 do
|
||||
raise ArgumentError, "expected rank to be 1"
|
||||
end
|
||||
|
||||
if Nx.type(tensor) != {:f, 32} do
|
||||
raise ArgumentError, "expected type to be :f32"
|
||||
end
|
||||
|
||||
bin = tensor |> Nx.to_binary()
|
||||
from_binary(<<bin::binary>>)
|
||||
end
|
||||
end
|
||||
|
||||
@doc """
|
||||
Creates a new vector from its binary representation
|
||||
"""
|
||||
def from_binary(binary) when is_binary(binary) do
|
||||
%SqliteVec.Float32{data: binary}
|
||||
end
|
||||
|
||||
@doc """
|
||||
Converts the vector to its binary representation
|
||||
"""
|
||||
def to_binary(vector) when is_struct(vector, SqliteVec.Float32) do
|
||||
vector.data
|
||||
end
|
||||
|
||||
@doc """
|
||||
Converts the vector to a list
|
||||
"""
|
||||
def to_list(vector) when is_struct(vector, SqliteVec.Float32) do
|
||||
<<bin::binary>> = vector.data
|
||||
|
||||
for <<v::float-32-native <- bin>>, do: v
|
||||
end
|
||||
|
||||
if Code.ensure_loaded?(Nx) do
|
||||
@doc """
|
||||
Converts the vector to a tensor
|
||||
"""
|
||||
def to_tensor(vector) when is_struct(vector, SqliteVec.Float32) do
|
||||
<<bin::binary>> = vector.data
|
||||
Nx.from_binary(bin, :f32)
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
defimpl Inspect, for: SqliteVec.Float32 do
|
||||
import Inspect.Algebra
|
||||
|
||||
def inspect(vector, opts) do
|
||||
concat(["vec_f32('", Inspect.List.inspect(SqliteVec.Float32.to_list(vector), opts), "')"])
|
||||
end
|
||||
end
|
||||
Reference in New Issue
Block a user