Skip to content

Commit 0eda1eb

Browse files
MageChiuMageChiu
authored
add ivf rabitq for command #553 (#554)
Signed-off-by: zhaopeng.charles <[email protected]>
1 parent 3583f4e commit 0eda1eb

File tree

1 file changed

+221
-0
lines changed
  • vectordb_bench/backend/clients/milvus

1 file changed

+221
-0
lines changed

‎vectordb_bench/backend/clients/milvus/cli.py‎

Lines changed: 221 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,162 @@ def MilvusHNSW(**parameters: Unpack[MilvusHNSWTypedDict]):
109109
)
110110

111111

112+
class MilvusRefineTypedDict(TypedDict):
113+
refine: Annotated[
114+
bool,
115+
click.option(
116+
"--refine",
117+
type=bool,
118+
required=True,
119+
help="Whether refined data is reserved during index building.",
120+
),
121+
]
122+
refine_type: Annotated[
123+
str | None,
124+
click.option(
125+
"--refine-type",
126+
type=click.Choice(["SQ6", "SQ8", "BF16", "FP16", "FP32"], case_sensitive=False),
127+
help="The data type of the refine index to use. Supported values: SQ6,SQ8,BF16,FP16,FP32",
128+
required=True,
129+
),
130+
]
131+
refine_k: Annotated[
132+
float,
133+
click.option(
134+
"--refine-k",
135+
type=float,
136+
help="The magnification factor of refine compared to k.",
137+
required=True,
138+
),
139+
]
140+
141+
142+
class MilvusHNSWPQTypedDict(
143+
CommonTypedDict,
144+
MilvusTypedDict,
145+
MilvusHNSWTypedDict,
146+
MilvusRefineTypedDict
147+
):
148+
nbits: Annotated[
149+
int,
150+
click.option(
151+
"--nbits",
152+
type=int,
153+
required=True,
154+
)
155+
]
156+
157+
158+
@cli.command()
159+
@click_parameter_decorators_from_typed_dict(MilvusHNSWPQTypedDict)
160+
def MilvusHNSWPQ(**parameters: Unpack[MilvusHNSWPQTypedDict]):
161+
from .config import HNSWPQConfig, MilvusConfig
162+
163+
run(
164+
db=DBTYPE,
165+
db_config=MilvusConfig(
166+
db_label=parameters["db_label"],
167+
uri=SecretStr(parameters["uri"]),
168+
user=parameters["user_name"],
169+
password=SecretStr(parameters["password"]) if parameters["password"] else None,
170+
num_shards=int(parameters["num_shards"]),
171+
),
172+
db_case_config=HNSWPQConfig(
173+
M=parameters["m"],
174+
efConstruction=parameters["ef_construction"],
175+
ef=parameters["ef_search"],
176+
nbits=parameters["nbits"],
177+
refine=parameters["refine"],
178+
refine_type=parameters["refine_type"],
179+
refine_k=parameters["refine_k"],
180+
),
181+
**parameters,
182+
)
183+
184+
185+
class MilvusHNSWPRQTypedDict(
186+
CommonTypedDict,
187+
MilvusTypedDict,
188+
MilvusHNSWPQTypedDict,
189+
):
190+
nrq: Annotated[
191+
int,
192+
click.option(
193+
"--nrq",
194+
type=int,
195+
help="The number of residual subquantizers.",
196+
required=True,
197+
)
198+
]
199+
200+
201+
@cli.command()
202+
@click_parameter_decorators_from_typed_dict(MilvusHNSWPRQTypedDict)
203+
def MilvusHNSWPRQ(**parameters: Unpack[MilvusHNSWPRQTypedDict]):
204+
from .config import HNSWPRQConfig, MilvusConfig
205+
206+
run(
207+
db=DBTYPE,
208+
db_config=MilvusConfig(
209+
db_label=parameters["db_label"],
210+
uri=SecretStr(parameters["uri"]),
211+
user=parameters["user_name"],
212+
password=SecretStr(parameters["password"]) if parameters["password"] else None,
213+
num_shards=int(parameters["num_shards"]),
214+
),
215+
db_case_config=HNSWPRQConfig(
216+
M=parameters["m"],
217+
efConstruction=parameters["ef_construction"],
218+
ef=parameters["ef_search"],
219+
nbits=parameters["nbits"],
220+
refine=parameters["refine"],
221+
refine_type=parameters["refine_type"],
222+
refine_k=parameters["refine_k"],
223+
nrq=parameters["nrq"]
224+
),
225+
**parameters,
226+
)
227+
228+
229+
class MilvusHNSWSQTypedDict(CommonTypedDict, MilvusTypedDict, MilvusHNSWTypedDict, MilvusRefineTypedDict):
230+
sq_type: Annotated[
231+
str | None,
232+
click.option(
233+
"--sq-type",
234+
type=click.Choice(["SQ6", "SQ8", "BF16", "FP16", "FP32"], case_sensitive=False),
235+
help="Scalar quantizer type. Supported values: SQ6,SQ8,BF16,FP16,FP32",
236+
required=True,
237+
),
238+
]
239+
240+
241+
@cli.command()
242+
@click_parameter_decorators_from_typed_dict(MilvusHNSWSQTypedDict)
243+
def MilvusHNSWSQ(**parameters: Unpack[MilvusHNSWSQTypedDict]):
244+
from .config import HNSWSQConfig, MilvusConfig
245+
246+
run(
247+
db=DBTYPE,
248+
db_config=MilvusConfig(
249+
db_label=parameters["db_label"],
250+
uri=SecretStr(parameters["uri"]),
251+
user=parameters["user_name"],
252+
password=SecretStr(parameters["password"]) if parameters["password"] else None,
253+
num_shards=int(parameters["num_shards"]),
254+
),
255+
db_case_config=HNSWSQConfig(
256+
M=parameters["m"],
257+
efConstruction=parameters["ef_construction"],
258+
ef=parameters["ef_search"],
259+
sq_type=parameters["sq_type"],
260+
refine=parameters["refine"],
261+
refine_type=parameters["refine_type"],
262+
refine_k=parameters["refine_k"],
263+
),
264+
**parameters,
265+
)
266+
267+
112268
class MilvusIVFFlatTypedDict(CommonTypedDict, MilvusTypedDict, IVFFlatTypedDictN): ...
113269

114270

@@ -156,6 +312,71 @@ def MilvusIVFSQ8(**parameters: Unpack[MilvusIVFFlatTypedDict]):
156312
)
157313

158314

315+
class MilvusIVFRABITQTypedDict(CommonTypedDict, MilvusTypedDict, MilvusIVFFlatTypedDict):
316+
rbq_bits_query: Annotated[
317+
int,
318+
click.option(
319+
"--rbq-bits-query",
320+
type=int,
321+
help="The magnification factor of refine compared to k.",
322+
required=True,
323+
),
324+
]
325+
refine: Annotated[
326+
bool,
327+
click.option(
328+
"--refine",
329+
type=bool,
330+
required=True,
331+
help="Whether refined data is reserved during index building.",
332+
),
333+
]
334+
refine_type: Annotated[
335+
str | None,
336+
click.option(
337+
"--refine-type",
338+
type=click.Choice(["SQ6", "SQ8", "BF16", "FP16", "FP32"], case_sensitive=False),
339+
help="The data type of the refine index to use. Supported values: SQ6,SQ8,BF16,FP16,FP32",
340+
required=True,
341+
),
342+
]
343+
refine_k: Annotated[
344+
float,
345+
click.option(
346+
"--refine-k",
347+
type=float,
348+
help="The magnification factor of refine compared to k.",
349+
required=True,
350+
),
351+
]
352+
353+
354+
@cli.command()
355+
@click_parameter_decorators_from_typed_dict(MilvusIVFRABITQTypedDict)
356+
def MilvusIVFRabitQ(**parameters: Unpack[MilvusIVFRABITQTypedDict]):
357+
from .config import IVFRABITQConfig, MilvusConfig
358+
359+
run(
360+
db=DBTYPE,
361+
db_config=MilvusConfig(
362+
db_label=parameters["db_label"],
363+
uri=SecretStr(parameters["uri"]),
364+
user=parameters["user_name"],
365+
password=SecretStr(parameters["password"]) if parameters["password"] else None,
366+
num_shards=int(parameters["num_shards"]),
367+
),
368+
db_case_config=IVFRABITQConfig(
369+
nlist=parameters["nlist"],
370+
nprobe=parameters["nprobe"],
371+
rbq_bits_query=parameters["rbq_bits_query"],
372+
refine=parameters["refine"],
373+
refine_type=parameters["refine_type"],
374+
refine_k=parameters["refine_k"],
375+
),
376+
**parameters,
377+
)
378+
379+
159380
class MilvusDISKANNTypedDict(CommonTypedDict, MilvusTypedDict):
160381
search_list: Annotated[str, click.option("--search-list", type=int, required=True)]
161382

0 commit comments

Comments
 (0)