Class: RagEmbeddings::Embedding

Inherits:
Object
  • Object
show all
Defined in:
ext/rag_embeddings/embedding.c

Class Method Summary collapse

Instance Method Summary collapse

Class Method Details

.from_array(rb_array) ⇒ Object

Creates a new embedding from a Ruby array



37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
# File 'ext/rag_embeddings/embedding.c', line 37

static VALUE embedding_from_array(VALUE klass, VALUE rb_array) {
  Check_Type(rb_array, T_ARRAY);           // Ensure argument is a Ruby array

  long array_len = RARRAY_LEN(rb_array);

  // Validate array length fits in uint16_t (max 65535 dimensions)
  if (array_len > UINT16_MAX) {
    rb_raise(rb_eArgError, "Array too large: maximum %d dimensions allowed", UINT16_MAX);
  }

  // Prevent zero-length embeddings
  if (array_len == 0) {
    rb_raise(rb_eArgError, "Cannot create embedding from empty array");
  }

  uint16_t dim = (uint16_t)array_len;

  // Allocate memory for struct + array of floats
  embedding_t *ptr = xmalloc(sizeof(embedding_t) + dim * sizeof(float));
  ptr->dim = dim;

  // Copy values from Ruby array to our C array
  // Using RARRAY_CONST_PTR for better performance when available
  const VALUE *array_ptr = RARRAY_CONST_PTR(rb_array);
  for (uint16_t i = 0; i < dim; ++i) {
    VALUE val = array_ptr[i];

    // Ensure the value is numeric
    if (!RB_FLOAT_TYPE_P(val) && !RB_INTEGER_TYPE_P(val)) {
      xfree(ptr);  // Clean up allocated memory before raising exception
      rb_raise(rb_eTypeError, "Array element at index %d is not numeric", i);
    }

    ptr->values[i] = (float)NUM2DBL(val);
  }

  // Wrap our C struct in a Ruby object
  VALUE obj = TypedData_Wrap_Struct(klass, &embedding_type, ptr);
  return obj;
}

Instance Method Details

#cosine_similarity(other) ⇒ Object

Calculate cosine similarity between two embeddings using optimized algorithm



107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
# File 'ext/rag_embeddings/embedding.c', line 107

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: %d vs %d", a->dim, b->dim);
  }

  // Use double precision for intermediate calculations to reduce accumulation errors
  double dot = 0.0, norm_a = 0.0, norm_b = 0.0;

  // Calculate dot product and vector magnitudes in a single loop
  // This is more cache-friendly than separate loops
  const float *va = a->values;
  const float *vb = b->values;

  for (uint16_t i = 0; i < a->dim; ++i) {
    float ai = va[i];
    float bi = vb[i];

    dot += (double)ai * bi;          // Dot product
    norm_a += (double)ai * ai;       // Square of magnitude for vector a
    norm_b += (double)bi * bi;       // Square of magnitude for vector b
  }

  // Check for zero vectors to avoid division by zero
  if (norm_a == 0.0 || norm_b == 0.0) {
    return DBL2NUM(0.0);  // Return 0 similarity for zero vectors
  }

  // Apply cosine similarity formula: dot(a,b)/(|a|*|b|)
  // Using sqrt for better numerical stability
  double magnitude_product = sqrt(norm_a * norm_b);
  double similarity = dot / magnitude_product;

  // Clamp result to [-1, 1] to handle floating point precision errors
  if (similarity > 1.0) similarity = 1.0;
  if (similarity < -1.0) similarity = -1.0;

  return DBL2NUM(similarity);
}

#dimObject

Returns the dimension of the embedding



80
81
82
83
84
85
# File 'ext/rag_embeddings/embedding.c', line 80

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);
}

#magnitudeObject

Calculate the magnitude (L2 norm) of the embedding vector



154
155
156
157
158
159
160
161
162
163
164
165
166
167
# File 'ext/rag_embeddings/embedding.c', line 154

static VALUE embedding_magnitude(VALUE self) {
  embedding_t *ptr;
  TypedData_Get_Struct(self, embedding_t, &embedding_type, ptr);

  double sum_squares = 0.0;
  const float *values = ptr->values;

  for (uint16_t i = 0; i < ptr->dim; ++i) {
    float val = values[i];
    sum_squares += (double)val * val;
  }

  return DBL2NUM(sqrt(sum_squares));
}

#normalize!Object

Normalize the embedding vector in-place (destructive operation)



171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
# File 'ext/rag_embeddings/embedding.c', line 171

static VALUE embedding_normalize_bang(VALUE self) {
  embedding_t *ptr;
  TypedData_Get_Struct(self, embedding_t, &embedding_type, ptr);

  // Calculate magnitude
  double sum_squares = 0.0;
  float *values = ptr->values;

  for (uint16_t i = 0; i < ptr->dim; ++i) {
    float val = values[i];
    sum_squares += (double)val * val;
  }

  double magnitude = sqrt(sum_squares);

  // Avoid division by zero
  if (magnitude == 0.0) {
    rb_raise(rb_eZeroDivError, "Cannot normalize zero vector");
  }

  // Normalize each component
  float inv_magnitude = (float)(1.0 / magnitude);
  for (uint16_t i = 0; i < ptr->dim; ++i) {
    values[i] *= inv_magnitude;
  }

  return self;  // Return self for method chaining
}

#to_aObject

Converts the embedding back to a Ruby array



89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
# File 'ext/rag_embeddings/embedding.c', line 89

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_new_capa(ptr->dim);

  // Copy each float value to the Ruby array
  // Using rb_ary_store for better performance than rb_ary_push
  for (uint16_t i = 0; i < ptr->dim; ++i) {
    rb_ary_store(arr, i, DBL2NUM(ptr->values[i]));
  }

  return arr;
}