diff --git a/README.md b/README.md index eae542d..28cdfb6 100644 --- a/README.md +++ b/README.md @@ -99,9 +99,17 @@ puts "Most similar text: #{result.first[1]}, score: #{result.first[2]}" - Embedding provider: switch model/provider in engine.rb (Ollama, OpenAI, etc) - Database: set the SQLite file path as desired -## 🔢 Embeddings dimension +# 🔢 Embeddings dimension -The size of embeddings is dynamic and fits with what the LLM provides. +In the previous version the size of embeddings was dynamic. Now the size of embeddings is static to use the RUBY_TYPED_EMBEDDABLE flag + +- the size of embeddings is currently set to 3072 and it's defined in two places: + - c side is defined in `ext/rag_embeddings/embedding_config.h` + - ruby side is defined in `lib/rag_embeddings/config.rb` + +Remember to recompile the c extension after changing the size: + +`rake compile` ## 👷 Requirements diff --git a/ext/rag_embeddings/embedding.bundle b/ext/rag_embeddings/embedding.bundle index d9c9296..4058184 100755 Binary files a/ext/rag_embeddings/embedding.bundle and b/ext/rag_embeddings/embedding.bundle differ diff --git a/ext/rag_embeddings/embedding.bundle.dSYM/Contents/Resources/DWARF/embedding.bundle b/ext/rag_embeddings/embedding.bundle.dSYM/Contents/Resources/DWARF/embedding.bundle index 31c89bf..fdf0418 100644 Binary files a/ext/rag_embeddings/embedding.bundle.dSYM/Contents/Resources/DWARF/embedding.bundle and b/ext/rag_embeddings/embedding.bundle.dSYM/Contents/Resources/DWARF/embedding.bundle differ diff --git a/ext/rag_embeddings/embedding.c b/ext/rag_embeddings/embedding.c index 28ff75f..d4ceab0 100644 --- a/ext/rag_embeddings/embedding.c +++ b/ext/rag_embeddings/embedding.c @@ -2,116 +2,58 @@ #include // For integer types like uint16_t #include // For memory allocation functions #include // For math functions like sqrt +#include "embedding_config.h" // Import the configuration -// Main data structure for storing embeddings -// Flexible array member (values[]) allows variable length arrays typedef struct { - uint16_t dim; // Dimension of the embedding vector - float values[]; // Flexible array member to store the actual values + float values[EMBEDDING_DIMENSION]; } embedding_t; -// Callback for freeing memory when Ruby's GC collects our object -static void embedding_free(void *ptr) { - xfree(ptr); // Ruby's memory free function -} - -// Callback to report memory usage to Ruby's GC -static size_t embedding_memsize(const void *ptr) { - const embedding_t *emb = (const embedding_t *)ptr; - return emb ? sizeof(embedding_t) + emb->dim * sizeof(float) : 0; -} - -// Type information for Ruby's GC: -// Tells Ruby how to manage our C data structure static const rb_data_type_t embedding_type = { - "RagEmbeddings/Embedding", // Type name - {0, embedding_free, embedding_memsize,}, // Functions: mark, free, size - 0, 0, // Parent type, data - RUBY_TYPED_FREE_IMMEDIATELY // Flags + "RagEmbeddings/Embedding", + {0, 0, 0,}, + 0, 0, + RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_EMBEDDABLE }; -// Class method: RagEmbeddings::Embedding.from_array([1.0, 2.0, ...]) -// Creates a new embedding from a Ruby array static VALUE embedding_from_array(VALUE klass, VALUE rb_array) { - Check_Type(rb_array, T_ARRAY); // Ensure argument is a Ruby array - uint16_t dim = (uint16_t)RARRAY_LEN(rb_array); - - // Allocate memory for struct + array of floats - embedding_t *ptr = xmalloc(sizeof(embedding_t) + dim * sizeof(float)); - ptr->dim = dim; + Check_Type(rb_array, T_ARRAY); + if (RARRAY_LEN(rb_array) != EMBEDDING_DIMENSION) + rb_raise(rb_eArgError, "Wrong dimension, must be %d", EMBEDDING_DIMENSION); - // Copy values from Ruby array to our C array - for (int i = 0; i < dim; ++i) + embedding_t *ptr; + VALUE obj = TypedData_Make_Struct(klass, embedding_t, &embedding_type, ptr); + for (int i = 0; i < EMBEDDING_DIMENSION; ++i) ptr->values[i] = (float)NUM2DBL(rb_ary_entry(rb_array, i)); - - // Wrap our C struct in a Ruby object - VALUE obj = TypedData_Wrap_Struct(klass, &embedding_type, ptr); return obj; } -// Instance method: embedding.dim -// Returns the dimension of the embedding -static VALUE embedding_dim(VALUE self) { - embedding_t *ptr; - // Get the C struct from the Ruby object - TypedData_Get_Struct(self, embedding_t, &embedding_type, ptr); - return INT2NUM(ptr->dim); -} - -// Instance method: embedding.to_a -// Converts the embedding back to a Ruby array static VALUE embedding_to_a(VALUE self) { embedding_t *ptr; TypedData_Get_Struct(self, embedding_t, &embedding_type, ptr); - - // Create a new Ruby array with pre-allocated capacity - VALUE arr = rb_ary_new2(ptr->dim); - - // Copy each float value to the Ruby array - for (int i = 0; i < ptr->dim; ++i) + VALUE arr = rb_ary_new2(EMBEDDING_DIMENSION); + for (int i = 0; i < EMBEDDING_DIMENSION; ++i) rb_ary_push(arr, DBL2NUM(ptr->values[i])); - return arr; } -// Instance method: embedding.cosine_similarity(other_embedding) -// Calculate cosine similarity between two embeddings static VALUE embedding_cosine_similarity(VALUE self, VALUE other) { embedding_t *a, *b; - // Get C structs for both embeddings TypedData_Get_Struct(self, embedding_t, &embedding_type, a); TypedData_Get_Struct(other, embedding_t, &embedding_type, b); - // Ensure dimensions match - if (a->dim != b->dim) - rb_raise(rb_eArgError, "Dimension mismatch"); - float dot = 0.0f, norm_a = 0.0f, norm_b = 0.0f; - - // Calculate dot product and vector magnitudes - for (int i = 0; i < a->dim; ++i) { - dot += a->values[i] * b->values[i]; // Dot product - norm_a += a->values[i] * a->values[i]; // Square of magnitude for vector a - norm_b += b->values[i] * b->values[i]; // Square of magnitude for vector b + for (int i = 0; i < EMBEDDING_DIMENSION; ++i) { + dot += a->values[i] * b->values[i]; + norm_a += a->values[i] * a->values[i]; + norm_b += b->values[i] * b->values[i]; } - - // Apply cosine similarity formula: dot(a,b)/(|a|*|b|) - // Small epsilon (1e-8) added to prevent division by zero return DBL2NUM(dot / (sqrt(norm_a) * sqrt(norm_b) + 1e-8)); } -// Ruby extension initialization function -// This function is called when the extension is loaded void Init_embedding(void) { - // Define module and class VALUE mRag = rb_define_module("RagEmbeddings"); VALUE cEmbedding = rb_define_class_under(mRag, "Embedding", rb_cObject); - - // Register class methods rb_define_singleton_method(cEmbedding, "from_array", embedding_from_array, 1); - - // Register instance methods - rb_define_method(cEmbedding, "dim", embedding_dim, 0); rb_define_method(cEmbedding, "to_a", embedding_to_a, 0); rb_define_method(cEmbedding, "cosine_similarity", embedding_cosine_similarity, 1); } \ No newline at end of file diff --git a/ext/rag_embeddings/embedding_config.h b/ext/rag_embeddings/embedding_config.h new file mode 100644 index 0000000..7a6c575 --- /dev/null +++ b/ext/rag_embeddings/embedding_config.h @@ -0,0 +1 @@ +#define EMBEDDING_DIMENSION 3072 // <--- this must be the same as what is set in ruby lib/rag_embeddings/config.rb \ No newline at end of file diff --git a/lib/rag_embeddings.rb b/lib/rag_embeddings.rb index 65f186d..f002530 100644 --- a/lib/rag_embeddings.rb +++ b/lib/rag_embeddings.rb @@ -1,3 +1,4 @@ +require_relative "rag_embeddings/config" require_relative "rag_embeddings/version" require_relative "rag_embeddings/engine" require_relative "rag_embeddings/database" diff --git a/lib/rag_embeddings/config.rb b/lib/rag_embeddings/config.rb new file mode 100644 index 0000000..cdec074 --- /dev/null +++ b/lib/rag_embeddings/config.rb @@ -0,0 +1,3 @@ +module RagEmbeddings + EMBEDDING_DIMENSION = 3072 # <--- this must be the same as what is set in C! ext/rag_embeddings/embedding_config.h +end \ No newline at end of file diff --git a/lib/rag_embeddings/database.rb b/lib/rag_embeddings/database.rb index baa2ddc..30f8f64 100644 --- a/lib/rag_embeddings/database.rb +++ b/lib/rag_embeddings/database.rb @@ -27,9 +27,13 @@ def all # "Raw" search: returns the N texts most similar to the query def top_k_similar(query_text, k: 5) query_embedding = RagEmbeddings.embed(query_text) + raise "Wrong embedding size #{query_embedding.size}, #{RagEmbeddings::EMBEDDING_DIMENSION} was expected! Change the configuration." unless query_embedding.size == RagEmbeddings::EMBEDDING_DIMENSION + query_obj = RagEmbeddings::Embedding.from_array(query_embedding) all.map do |id, content, emb| + raise "Wrong embedding size #{query_embedding.size}, #{RagEmbeddings::EMBEDDING_DIMENSION} was expected! Change the configuration." unless emb.size == RagEmbeddings::EMBEDDING_DIMENSION + emb_obj = RagEmbeddings::Embedding.from_array(emb) similarity = emb_obj.cosine_similarity(query_obj) [id, content, similarity] diff --git a/spec/performance_spec.rb b/spec/performance_spec.rb index 8445570..9435d38 100644 --- a/spec/performance_spec.rb +++ b/spec/performance_spec.rb @@ -6,7 +6,7 @@ let(:text1) { "Performance test one" } let(:text2) { "Performance test two" } let(:n) { 10_000 } - let(:embedding_size) { 3072 } + let(:embedding_size) { RagEmbeddings::EMBEDDING_DIMENSION } let(:emb1) { Array.new(embedding_size) { rand } } let(:emb2) { Array.new(embedding_size) { rand } }