@@ -13,8 +13,8 @@ use serde::{Deserialize, Serialize};
1313use crate :: {
1414 comparator:: {
1515 Centroid16Comparator , Centroid16Comparator1024 , Centroid32Comparator , Centroid4Comparator ,
16- Centroid8Comparator , Disk1024Comparator , DiskOpenAIComparator , OpenAIComparator ,
17- Quantized16Comparator , Quantized16Comparator1024 , Quantized32Comparator ,
16+ Centroid8Comparator , Disk1024Comparator , DiskOpenAIComparator , Memory1024Comparator ,
17+ OpenAIComparator , Quantized16Comparator , Quantized16Comparator1024 , Quantized32Comparator ,
1818 Quantized4Comparator , Quantized8Comparator ,
1919 } ,
2020 openai:: Model ,
@@ -28,6 +28,7 @@ use crate::{
2828} ;
2929
3030pub type OpenAIHnsw = Hnsw < OpenAIComparator > ;
31+ pub type Memory1024Hnsw = Hnsw < Memory1024Comparator > ;
3132
3233#[ derive( Serialize , Deserialize ) ]
3334pub enum HnswConfigurationType {
@@ -36,6 +37,7 @@ pub enum HnswConfigurationType {
3637 SmallQuantizedOpenAi8 ,
3738 SmallQuantizedOpenAi4 ,
3839 UnquantizedOpenAi ,
40+ Unquantized1024 ,
3941 Quantized1024 ,
4042}
4143
@@ -92,6 +94,7 @@ pub enum HnswConfiguration {
9294 DiskOpenAIComparator ,
9395 > ,
9496 ) ,
97+ Unquantized1024 ( Model , Memory1024Hnsw ) ,
9598 UnquantizedOpenAi ( Model , OpenAIHnsw ) ,
9699 Quantized1024By16 (
97100 Model ,
@@ -115,6 +118,7 @@ impl HnswConfiguration {
115118 HnswConfiguration :: SmallQuantizedOpenAi4 ( _, q) => Some ( q. quantization_statistics ( ) ) ,
116119 HnswConfiguration :: Quantized1024By16 ( _, q) => Some ( q. quantization_statistics ( ) ) ,
117120 HnswConfiguration :: UnquantizedOpenAi ( _, _) => None ,
121+ HnswConfiguration :: Unquantized1024 ( _, _) => None ,
118122 }
119123 }
120124
@@ -138,6 +142,9 @@ impl HnswConfiguration {
138142 HnswConfiguration :: Quantized1024By16 ( model, _) => {
139143 ( HnswConfigurationType :: Quantized1024 , model)
140144 }
145+ HnswConfiguration :: Unquantized1024 ( model, _) => {
146+ ( HnswConfigurationType :: Unquantized1024 , model)
147+ }
141148 } ;
142149 let version = 1 ;
143150
@@ -156,6 +163,7 @@ impl HnswConfiguration {
156163 HnswConfiguration :: SmallQuantizedOpenAi8 ( m, _) => * m,
157164 HnswConfiguration :: SmallQuantizedOpenAi4 ( m, _) => * m,
158165 HnswConfiguration :: Quantized1024By16 ( m, _) => * m,
166+ HnswConfiguration :: Unquantized1024 ( m, _) => * m,
159167 }
160168 }
161169
@@ -168,6 +176,7 @@ impl HnswConfiguration {
168176 HnswConfiguration :: SmallQuantizedOpenAi8 ( _model, q) => q. vector_count ( ) ,
169177 HnswConfiguration :: SmallQuantizedOpenAi4 ( _model, q) => q. vector_count ( ) ,
170178 HnswConfiguration :: Quantized1024By16 ( _, q) => q. vector_count ( ) ,
179+ HnswConfiguration :: Unquantized1024 ( _, q) => q. vector_count ( ) ,
171180 }
172181 }
173182
@@ -182,7 +191,8 @@ impl HnswConfiguration {
182191 HnswConfiguration :: UnquantizedOpenAi ( _model, h) => h. search ( v, search_parameters) ,
183192 HnswConfiguration :: SmallQuantizedOpenAi8 ( _, q) => q. search ( v, search_parameters) ,
184193 HnswConfiguration :: SmallQuantizedOpenAi4 ( _, q) => q. search ( v, search_parameters) ,
185- HnswConfiguration :: Quantized1024By16 ( _, _q) => {
194+ HnswConfiguration :: Quantized1024By16 ( _, _)
195+ | HnswConfiguration :: Unquantized1024 ( _, _) => {
186196 panic ! ( ) ;
187197 }
188198 }
@@ -195,9 +205,8 @@ impl HnswConfiguration {
195205 ) -> Vec < ( VectorId , f32 ) > {
196206 match self {
197207 HnswConfiguration :: Quantized1024By16 ( _, q) => q. search ( v, search_parameters) ,
198- _ => {
199- panic ! ( ) ;
200- }
208+ HnswConfiguration :: Unquantized1024 ( _, h) => h. search ( v, search_parameters) ,
209+ _ => panic ! ( ) ,
201210 }
202211 }
203212
@@ -225,6 +234,7 @@ impl HnswConfiguration {
225234 HnswConfiguration :: Quantized1024By16 ( _, q) => {
226235 q. improve_index ( build_parameters, progress)
227236 }
237+ HnswConfiguration :: Unquantized1024 ( _, h) => h. improve_index ( build_parameters, progress) ,
228238 }
229239 }
230240
@@ -253,6 +263,9 @@ impl HnswConfiguration {
253263 HnswConfiguration :: Quantized1024By16 ( _, q) => {
254264 q. improve_index_at ( layer, build_parameters, progress)
255265 }
266+ HnswConfiguration :: Unquantized1024 ( _, h) => {
267+ h. improve_index_at ( layer, build_parameters, progress)
268+ }
256269 }
257270 }
258271
@@ -280,6 +293,9 @@ impl HnswConfiguration {
280293 HnswConfiguration :: Quantized1024By16 ( _, q) => {
281294 q. improve_neighbors ( optimization_parameters, last_recall)
282295 }
296+ HnswConfiguration :: Unquantized1024 ( _, h) => {
297+ h. improve_neighbors ( optimization_parameters, last_recall)
298+ }
283299 }
284300 }
285301
@@ -308,6 +324,9 @@ impl HnswConfiguration {
308324 HnswConfiguration :: Quantized1024By16 ( _, q) => {
309325 q. promote_at_layer ( layer_from_top, build_parameters, & mut progress)
310326 }
327+ HnswConfiguration :: Unquantized1024 ( _, h) => {
328+ h. promote_at_layer ( layer_from_top, build_parameters, & mut progress)
329+ }
311330 }
312331 }
313332
@@ -319,6 +338,7 @@ impl HnswConfiguration {
319338 HnswConfiguration :: SmallQuantizedOpenAi8 ( _model, q) => q. zero_neighborhood_size ( ) ,
320339 HnswConfiguration :: SmallQuantizedOpenAi4 ( _model, q) => q. zero_neighborhood_size ( ) ,
321340 HnswConfiguration :: Quantized1024By16 ( _model, q) => q. zero_neighborhood_size ( ) ,
341+ HnswConfiguration :: Unquantized1024 ( _, h) => h. zero_neighborhood_size ( ) ,
322342 }
323343 }
324344 pub fn threshold_nn (
@@ -346,7 +366,12 @@ impl HnswConfiguration {
346366 }
347367 HnswConfiguration :: Quantized1024By16 ( _, q) => {
348368 Either :: Right ( Either :: Right ( Either :: Right ( Either :: Right ( Either :: Right (
349- q. threshold_nn ( threshold, search_parameters) ,
369+ Either :: Left ( q. threshold_nn ( threshold, search_parameters) ) ,
370+ ) ) ) ) )
371+ }
372+ HnswConfiguration :: Unquantized1024 ( _, h) => {
373+ Either :: Right ( Either :: Right ( Either :: Right ( Either :: Right ( Either :: Right (
374+ Either :: Right ( h. threshold_nn ( threshold, search_parameters) ) ,
350375 ) ) ) ) )
351376 }
352377 }
@@ -372,6 +397,9 @@ impl HnswConfiguration {
372397 HnswConfiguration :: Quantized1024By16 ( _, q) => {
373398 q. stochastic_recall ( optimization_parameters)
374399 }
400+ HnswConfiguration :: Unquantized1024 ( _, h) => {
401+ h. stochastic_recall ( optimization_parameters)
402+ }
375403 }
376404 }
377405
@@ -387,6 +415,7 @@ impl HnswConfiguration {
387415 }
388416 HnswConfiguration :: UnquantizedOpenAi ( _, h) => h. build_parameters ,
389417 HnswConfiguration :: Quantized1024By16 ( _, q) => q. build_parameters_for_improve_index ( ) ,
418+ HnswConfiguration :: Unquantized1024 ( _, h) => h. build_parameters ,
390419 }
391420 }
392421
@@ -398,6 +427,7 @@ impl HnswConfiguration {
398427 | HnswConfiguration :: SmallQuantizedOpenAi4 ( _, _)
399428 | HnswConfiguration :: UnquantizedOpenAi ( _, _) => 1536 ,
400429 HnswConfiguration :: Quantized1024By16 ( _, _) => 1024 ,
430+ HnswConfiguration :: Unquantized1024 ( _, _) => 1024 ,
401431 }
402432 }
403433}
@@ -416,6 +446,7 @@ impl Serializable for HnswConfiguration {
416446 HnswConfiguration :: SmallQuantizedOpenAi8 ( _, qhnsw) => qhnsw. serialize ( & path) ?,
417447 HnswConfiguration :: SmallQuantizedOpenAi4 ( _, qhnsw) => qhnsw. serialize ( & path) ?,
418448 HnswConfiguration :: Quantized1024By16 ( _, qhnsw) => qhnsw. serialize ( & path) ?,
449+ HnswConfiguration :: Unquantized1024 ( _, hnsw) => hnsw. serialize ( & path) ?,
419450 }
420451 let state_path: PathBuf = path. as_ref ( ) . join ( "state.json" ) ;
421452 let mut state_file = OpenOptions :: new ( )
@@ -473,6 +504,9 @@ impl Serializable for HnswConfiguration {
473504 QuantizedHnsw :: deserialize ( path, params) ?,
474505 )
475506 }
507+ HnswConfigurationType :: Unquantized1024 => {
508+ HnswConfiguration :: Unquantized1024 ( state. model , Hnsw :: deserialize ( path, params) ?)
509+ }
476510 } )
477511 }
478512}
0 commit comments