From 3f4c7804e8d444dccdf52a314c8f16fa943cadc7 Mon Sep 17 00:00:00 2001 From: Marco Mastrodonato Date: Sun, 8 Jun 2025 19:58:47 +0200 Subject: [PATCH] Now the size of embeddings is static and the RUBY_TYPED_EMBEDDABLE flag is used - the size of embeddings is static and 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` --- README.md | 12 +++ ext/rag_embeddings/embedding.bundle | Bin 51920 -> 51504 bytes .../Contents/Resources/DWARF/embedding.bundle | Bin 16192 -> 14996 bytes ext/rag_embeddings/embedding.c | 94 ++++-------------- ext/rag_embeddings/embedding_config.h | 1 + lib/rag_embeddings.rb | 1 + lib/rag_embeddings/config.rb | 3 + lib/rag_embeddings/database.rb | 4 + spec/performance_spec.rb | 6 +- 9 files changed, 42 insertions(+), 79 deletions(-) create mode 100644 ext/rag_embeddings/embedding_config.h create mode 100644 lib/rag_embeddings/config.rb diff --git a/README.md b/README.md index ef13d54..28cdfb6 100644 --- a/README.md +++ b/README.md @@ -99,6 +99,18 @@ 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 + +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 - Ruby >= 3.3 diff --git a/ext/rag_embeddings/embedding.bundle b/ext/rag_embeddings/embedding.bundle index d9c9296a6bffc2735fb84c3dabe100fb299c34df..4058184511349195e6ca7fb640aef1e759b4ba77 100755 GIT binary patch delta 2067 zcmaJ?4Qx|Y6h800uItZ)b}OR@D_zIN1{>=J%7A2aLkUK7Zv06^){ahEvvmvGOkikB zPyq$$>m3O}Vum7a9mK+bAR=_lpd%^_L_(s$E)ra(F-`_D!l>uI-WqI-H~HQ<-}$-c zdvCXQ{dO}lq{-MLBi{)5 zyje*e2&UDvKoJBeiPcL_3vvqw;OPM{VuVNN5go|$(MEa!%otRo{}O_($qd*AQxFzr?M%BC7P32PcVH;Hf%|IVrxA z#WfjXL`d@sm^%{BwXtZP;I0s_T<6{*gX0dB#kWKT zCmg(GRH?Q27G}V3j{wURGRIuZjZWpv!v~O?s1hrP9K2h#9d{f-mFXT36;{iBov28p z(xhgTOb1AxC&@Aj_<0nbEmea8#6HttuJ|VA@2&a@`?=1sR=#s=D~Or&6QaPin4_N% zuxADOR45f@eZA!j1Wg79#s)AJ5i~kK{`uQZq@uEM7F0&AC+94^BXIsXxK~>Zhyy6J z982y>_4LJb{!-)$BZ!?Q5aAXSs&Ed6mPcf1-38f`Lix#M?gH&|p?tDaVN2J-uuejG%^P9Wh6dMq zqup6w>uPc~xbltlElthFnp$JFo!Hcq$Oq~fd_X}?tEZ&%pp$-0s1|_30Wty6txZXu z(``(7mg85d@RXO5ezmMymm;IVcD2pU?utfcH#57P**0d|nLW(xDQ0__?PvB!W(S!a zAw`zyIO_gve0bE!+Dv_=30@`RgI)4Xwz|58b>wU?gi3*Y+LK-3<*soK$V2tx5RGz- z<66l0Gs!Vt!T4du>nKNGZ;9nMQ;q`g)(TCW@jfzpEFnF^3H&?cuK@>NeoK6zS)7;j zgr>UVqo2fI`7|tub97-E;u_nK@)r~GvyjK7Qt7lm&e4_n@uW+Oa@zlR9{cA*{N3og z(hcIAHc$aF>foY=ZV=Dck=4Ds-SnVhkktVAByRbKA56%fOvqnN$VU?R?+N_DzqwnI z_qYIkjNsRdEqagBE7*WAA~F#sL>7X+UN&Mn0&g1W>YjmiDk29l6Y>B3Y5%bS$-B@R zKXM%xJ|fm|F*y(}Aj9D=NTBa&d01oQAqD@nuYN5cBgCoteO^=y{;D4 zZR=XQ*n8d;`pG+bqwlkh=RfSd_}687!{=AFzI`>V@)&vHqSg5FA1!xpn;)ecy3Q?A z?$|g~V#<4E$a~CJaj0aCXQU`96?GM_|8w_1-u`(v*PNd?^Kjw0y_dT`DLQs!!R^AM u%iO!h)1s_uvw7KXyS5CD5PPXV3YauY2yf zA8)%_mBF2g-!JzmTIQzV60gQjP;WCmm+lxnP`87>S%2@(Z?0CWRCqv(^E(rm*Xx-P7qiR~C`eh|gwp~&38RcBxOdzg@SVF-))RA0>D zz?ec>GVdtj-hj?Rj;6V=Cz&mX_jT(l9WD9{i6pEND}C-pz$*ygTR;ff23cR7>}YD2 zXmJ*j6mQvdh=M7JufaA1KVFy*GP=?fK;egF%Y_le+R^*(cpEx8q2H{8Lg6124$cP zoO8_e_W)b7wzjHCq*bPn#&s7_mGF;RP^$Ut?$;gm{BECp?Jxo0J@p7SQ>BRQm2$*(i)1}sy25d2IjF>9vbju(t%XcBi zo@5YjQ*zM7d=~p+oc0*~#G4)0$T;!m?HY7#cGy~o_VazRU@m@e_+eySmNygEpi48B zAV{{I*5&MwycYU}Kd}PqI$sm+7)s6(HD+a*$4d=0P8XHd)O)Jk z&U%m6Q{}4jG&J4pzF;DfmV3PJ>UvL2b&}Up?{zhlRuJ>5^W{BOwo9FpVWuQ9qdKBa zBB@Fu8P{LwRLF!9OjAKHnUoNh_kE$pES^=u-?Xr?L%iR9=4!c9lHin|weOr-tomxj z8~%*oFACl!_-?^F1wSZwpWuUnpAr1L;KPFJuPZ8o9A9fG zNTK3Ly>K=)_`d1sI#(?_80@*Dt|H3etlL!5S`x6URPCBO)2@=v)y<2 zI;Eg=SL1GlcZDbyJ>uKBV~P%`NZ?1&H$FB9+#v8~fiDWo1U4G@j=+A=&;uO9H{PMA zqkddYnhs+lyVN~%-q%27A&dSSg?BL=Cjf6Ci8?=Fi+VDoPUh;d`HWi2U`#~<14E=# zP*6{WvjG0sBzp$1u{OR|9&ccXdOWV0J~-y_;daFP$Y=&{4o_fU9>5DQzTh151mSjG zfYChOOdfAwh5IdZpxr;{c2-!oIc;7+&kvO2wPZ@p)W zn?2HhB=iol$fHZA40I*QG!1{x((YCIP^Ug7f-*`rRhLz{=~#mOKw5nH(Um8L)nE5z zZ+ywAr562L^mzPWL)xC^oull7qs2*IU7s6U_0P6HKJ4lJ{=KRnBiC&inZb@n>^ife z{+VjeyJp>K)~9ifwxZJ$#q3Lg)H1)sz(SdyKbi41*%pc#X hsC_~GezbAziFHcAn*VwE_1{L!1wEH5{cCtH{{uGcjIRIy 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 31c89bf28aeef8e579c94c46de73f2ab75ce1a46..fdf04184d5fb768db81792fcc89ecd3fe739ef80 100644 GIT binary patch delta 4394 zcmai1d303O8UOCvGLy+NlgTpKGbAJfB!MIo5&}rTuw+082n1Bf4vA)KNI(%7NZCZG z!G|p(6|rc=gBA?p5fNxli&hIwZK`OsqIiTO6_rI%Tif6F-V3CEbL+Vb_SL+DX*8tznteCvAIwiv55MI{bF& z!^BiEL-&eQy+@3)Q>dwL)m?I&eioW`OH-}KOT8joKP9~-8uf!xoPI!7gj0V>9s&7d zxk~CMeUdUo+Mb|qQErjiCH+lhN>uyXX24M=$oS5tM}>h+E8RkQSB&EvY%Z>ri=lU~wYT9&?9^F%Fv-_-O=CS5$- zm_(nJ7M1^@srp(mX?~7jN()JgO8>-E?Xo73o^P1iVLE=!G`X!_#|sS8iIB9I{`02k zgf&)w5s}n8(dDo`;J33k7YYi<62X$*Vj4o z_1`&jdZS#IB6O#%HrR{-RkG1sgcU;k5E0IV=WZ%ejk%76JQa8lkVUHGP*Wn_0Z(`T z_+;~jD}i;Fk8b$t;SUy7a`NKj=;S2_zCAQ_<2AA9ySm1yBRN#O=jCf5e?0lz*s|_< zkN3^YSplW=Vz51WR^&I`qu=@MwSoz{MJc~v;jkX(60(bR}g`7Qk`Z%k0J$$e|O)%2MPsUh;# zH32npqaxt;fMOe`hG_w{u3P@-{s7EFH8~-mCL>~OK#dNl*;+Yb2ekE?BY=>qChwJP zJRppp~N}6oHtjeKR)yONFTy;cE(d6_1rq->hwcTorw(dq4_+Epj z?H1ek*VHgumzoHtY)xu(sZpEIFQCQ;)SQ6o3#f&*PM$}oBOp1kX|hpzEF^yBEitef zeUe_2?hEz^oQSa~8t*p#389cjRA6T)v_RC6kQMC03T5uILiv@b*sVM(_6x}S#UV>^ zNSt8B3DIvVj4mEATk-4?qE6W*jGnW(XFn7q5)>*HOAH$Hi+U2OLP_xCxv-r@sc5$$ zgOac-_MoRF;gJtrsKO8sLmY z>n(e`Mf)w5;w3JYxCk)(J8qN2rCmQ@XDq=qVLo_1( zB3g+O#)r;lW%w&_poCjOTMP+J(A0&qOo6+_0>#IGj-kd7cuq`R9ip7G__b5EWL2cbOt@meDjcvJq6Y zET9yl7SBLhX!4%J;Ac^)vHr)v-$7pX$#oh_auS-^G{lISw*`}RVdgUKM#|$mBZ{S< zbFGKs_d|5&9YGt-$*o7nf;(|<@+L9oVq}zXmt!Hv4>xmJ6gI|?I3X%n?8qIIEgNWP zi~=HBswj))%(KufwD_;VS5sDy4IUXyIr)q{!ob&1HXl7t0-9Qj_kpjo_|f3&P5<0j zK8i(1H0Nfbaa_RlTtG8jPCoXAj*D8%DdT%c-*~3=u6fghxvu%;(62$A3I?ZRF0lWmiWQp#>QktD zN4ir-^M`H`f`{J>SuBT522W>zi7aR&~F8jGy2ggOd(5GEEWN}(Bc`>jqot~r3{pe&ZJjtYJM@14N5S*B^ezEcf$PA( zfiN6b7mx(x0Jj2*fEu6#lKeK{yb9EAV51T}f=^UnIaUgA3r-n*Ro3 z-w?JkeOB21*yA7nB=X%}Z)@SmovZGVR=-sd?*B!{;lA%yIA3?I*ju;q>DusZkcPV&EANS%*XEVR#Md9jjRk%7)$cI=JKde6rI*;gX1!z^W1e%(NUk{GK)L2I3sh!Wt2C*866!TIy%n#eRc0mnxJz$C;s!^ z$vx*+^?hG`^;OkZRky0^R=@c7KYx+Q7-x7g06XwO8ga%`geV`l&wpy~hEaQ{7A|uOTOabuNOgVEhmJWIk;vXO`0A7W%u6H-}bb}fH)QuMyyn@IANpw9K z+?+Jh4E|_)he#REQ<4!=NOc;0WVU^j2BGVT?ZJoQ zG3@CB!SG-(4#{|XJ>+CO&l!;jQ|+;1>gP`+S z_3DqMbSwBTw?|{qdQ!TWsSp|?&`8yFe=M39vpC*fex0GW%7{dmYA;1BNE*Lm?MzAC zZqjurvJ^74og~)vuuo5nM0#UkeJ~j}+nJJKyFsV=Keb*oHK(Mg>w3VK^o9z zw;MFgYwJn1XUVLX+DSgqb+4ywbyK6q=uO6+Fg2*N=rQ9(#tT5aKul63lbwf95bCq8 zzH_CJiQpxZEIm4|n_=+OKB}KN43R$&>>cXY!=Y%9G``ASv(>0l()#MhvgeEYqW!_d zG4`4d7=#;)%~^Y>9M}PWBT{N+eZF_rvDDq zem#3Mp0V`A$Jg)fcZ~dvX8rD7^!4mTeG!c2WBvEH_l*2rbAJ5oq;c!vzcxRaI_ZH) z4@`Ps(gTwonDoH^Ll4llDjR#3OU(Tm{k|gJYn}njHTk5;FE;uAHhKDaAo82PlLIEd z$K;2JhpzQ|Yd7&ovwoLc4L%3uS-;z{g9%1@gMk2k^6Gu@SVZ^5OA>OdSKE6lOCA#z@!HzJuvBkNe@hVVA2DV9+>pNqzC>ld0+*;3n=*Ne+PbpbIw1@ z&Y}nmebECkLwtw8M+4RjTnx|$9*NJeIh5&BN=R}2?-sqg?xxdc-2VGt)!g-`1zpEc z&Sw^t^xbphf!Vu0=BIz^A86i0!h-wEiY5{&Z+%R7caah%YMUktN(pYSEq)< z3x3DW^)J75?^W4+!4>Zg-1_X>8%Vh1n%WC0L$}8|Us$`ZVb8%=ZzJJ`D z5>DbTzWZ+%S8l)Q`j6R%?M1^UvWrMqdGj5^mu)ZT{_X4DSyx`qPkWJsw-)EtexBUu zU-duB_ifnpOYCNO;ogjoYLn-fK@UcYXTAN746Ly3IN=k(qRuamAVO zsIvSKWtzHQk<^dv!}dGvcPiY0fJt|4+b^q1c80XyW=|^9r8m#ruQ;Xe^A;stI>1^~ z>G`cYWVBq=OMlKNcAKuoYOjo6!F%+!sP?Q}=IZv4i4XH)_iRwT;^$0TBm)n;q89z}pq@*EH zy-<~I*|uLg$n0;UzV<&VMaVuwwGHGV@k0{b2@x7i%~`!#nc6HWJw@$Ac`Zt1i&D^{ zRJ14y;Uu&x10`jvQsP6jhMbxqJ+y7Vl9AD&oB(Zii&9QCvdNey^i>ig2aoZg-Rhgv zYbYYw3|xq5+R?)mnH9@7@Y^^$3&VU=%w?9(I~boXvn2O3F3D^O?~s?=#(6+~fIq}3 z5yE_aikuIjU%rh$z~!6yDSMIp1O$xFkXbhm8bv^X3L!C-Xj{Bj@Q4{tmzh(_WqW03 zE$?GITV@)H(kfuB1IbRAWm0Lm$On1xW>a1&vktO1MP`jWAmU76CbL|+VB21~yp{^j zlUX@(mES0t*&xtRX>A@FGOv~l&y?A#HdOGC?Gr?d7s~9AjS})@c9^T*!xX?|J&&yxTTJ;qX8e>HzhuUT z&G?WR%Vs{i880uhzvx}hAliy>;H<HTW7G<=H(8j8JC*zpm_#+ zi5XjGvYJ`Hx6SxU)1GGLyBu zTQhjIodW)JTG}f9GrRp^aQsD9i2c;??QSS4ddFm!#*?`7>Y0>zJbqGW8h*Q)~l{Un0)TA~lxn zcokvJTL`j_kU*u+;x4&H{Rpupt&^sv*&UyQcIHxsDKwusb5Og4$~2|ll|(UlvU;|b z&$I%r7232SNh_AM5=EP?YBSQbQoA-YU7MAml{vKXOl@|SHYZz~o1<0aYR65{j(2G% zXxhA~TIDpYs!($mY1PHr{1UBZy0%~jlWNmsbvsH}m_9>VlxC9_%d)gYk!{pC>0653 zaUBTTUIgk75uC^z88WlmPudIyH|^|8X)$Gxb<f^%d%+>aPUQP*K`#*(X zdaM0CA!xI|Cp6pbABcom(i13&M><#;@nyx}PBH2EpgRpZuLE=!tN1yUauSBXd; z$z4bafyZiIb z2Io*as!)2mbCZbAa1M79uRAx3_)O=>C1l7(bIe9rX8Lm97bco>-xoG#Wz&0vM|#;J zl9%lU=QrrQ?}F}Pi>c?1Uk|%E0LDGu8a+Oc$4|3lxgWx#Jy<1K4eyyFv?^al-S#pku`yLosa9xu7F#J-rI@<|2ri^fJ&fgPxZ! zQ)(B2Znj~pH3c-{E}#)7oF`gXI4YyoxT%3f`IT5EirIp}oFQYB(EFm6Ei9%U;IXq= z_9EWV#j?roi+!YPnPf}inT6bVGtVla$;D#{=AuDcb_S}`#a#4~&g%o6WZG4TppwgH z7rGRwFrMekD=N$_nxZE1Fs|k>4>nU$Ig<)cb~qfEvqU-^CxcXKL3&&#I!h9N%9pjGgGxiNwcc>`1+G#nQ~P9%S-rt*P6T@$9l&)_+&Pv&WGBZ3^w1Di+0T= z1u_H`g;Xk_l20nfDXch$h~raKgq0I&kr{K-A!@)_j@5xo-ZLhO?t?i9# zyrYUMz3ug*dR^U}p0;L7v9qbUsi(u@gbOWZRejegi$4IWc8xM!>suSzN7Yt(dQ$4J zrmeAQrDe?1+SJ-^F>URQqZT|z`RLU4)&|dLT5J2-rba7MYkk*hHrgv8^yv_S`HRqZ zqte7cES_Y2;g~PUNJI49=!k|UUF#4Xm|){>TSPj2{mkEFT?-K-G1y0!JR-5cP&k+( zz^q9m@E3;w{(KSd4MhgSdS`<%zB}|jUm`i^OAfGD?^#$_^oCe0%vKBp{Tuae8Z(5;c6-`{4Gdi|55$!eJsgZOAL}&cQB!L}Q={%}YGc0m2sR93 zElat?vXN>>wKeXv(Eb0AF>*IYd|_PqHqu6PObb|nVst1{6X*@IKy0YjY$&e#i^>Ib zGy`?jX%dp|msH-<)?I_1XJ=u<5XOceN_y~q4eTNom+HrAN?DIn(O zn3lM2g9eZm+Nwu`!@A!$=p$ZqFG;M#!J))JiZ716n-fb8psRXWG!~EOX5>S2g`eZT zP$HO;UA#@iOy4ZXu}WF@cpwpfGr|Ui(YArE@h=62fvvzEfOdL!0Z#%4fP=uBz+vDo z0PQ^}a~hBdOaW*)tpJt)D*-RC0T=-2X!Cr4wszkG=qPnRKnIxzfH#2mfzN<6Z29Tf zvJ6-NoD8%8>wqAT0JZ~{0@ngR1bz%W3A_Nj2D}Y?1bhn6aitBL$ZTK=Pykc_mB13f z1L(jAa1O8yxB%D#dPa3ydpa0l=R@Hp@ya1b~Q$mp;PAQNx_g}`)RE`UXq)dRG-p1`t~ zl9jo=q24$Kue&k0DHx97uUhUEu}H)h4S4Zd>Bf5@;qGc&op87M{Ow&mRf)hx_fR74 z4n_Upp+L~>i$^AGBQcV2i;0|=uqo5b=ybsTpvspB^r60qBwj2DH%%J%f~wl8+R-+U zsFFn0Mk-zl!FbdccKc&-Y$Haee$~LljnOLvc8q!wVXWj6Gc6X2Z@`48)x_*1uxNyQ zVdFn%iHX@V+%HxL7}Bkk^qVl+>q~_E6EkG2xD&HsE_M?$VhmTzMBEOYm?gR-IuHpa zLw>ih%udW)FdmP^CuSn0Mzl@(#*D#du*Fv4g^qWrSRiSwqbCG9-`+nnI_!KlclZ|nmMcDdsc2!kQ z7jOx1HE;uP6L1?qb$SS(^Yo{HUjQ!wzX1*cZvuw_UM{(0be^pQF*;PLGI+Vxh!xq0 z)oDhYmVp?0c9+33E>4le$1rB=G`+#3TWl~cS0iAGaX`F`nk>%XnJUh^xHuP+{zx9N zIgpdu#`s4Rs^W+{O|hi5cVUlwg~K1nRuaX // 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 8caf430..e8b617e 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) { 768 } + let(:embedding_size) { RagEmbeddings::EMBEDDING_DIMENSION } let(:emb1) { Array.new(embedding_size) { rand } } let(:emb2) { Array.new(embedding_size) { rand } } @@ -31,8 +31,8 @@ puts "RSS: #{(`ps -o rss= -p #{Process.pid}`.to_i / 1024.0).round(2)} MB" # Weak expectations, mostly for sanity check - expect(creation_time).to be < 0.1 # Should be less than 100 milliseconds for 10_000 - expect(sim_time).to be < 0.1 # Should be less than 100 milliseconds for 10_000 + expect(creation_time).to be < 0.2 # Should be less than 100 milliseconds for 10_000 + expect(sim_time).to be < 0.2 # Should be less than 100 milliseconds for 10_000 end context "memory usage" do