Skip to content

Commit 40dd2ba

Browse files
egolearneregolearner
andauthored
cli support run LabelFilterPerformanceCase (#626)
* cli support LabelFilterPerformanceCase * fix format --------- Co-authored-by: egolearner <[email protected]>
1 parent a118380 commit 40dd2ba

File tree

1 file changed

+17
-3
lines changed

1 file changed

+17
-3
lines changed

vectordb_bench/cli/cli.py

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -188,6 +188,11 @@ def get_custom_case_config(parameters: dict) -> dict:
188188
"dataset_with_size_type": parameters["dataset_with_size_type"],
189189
"filter_rate": parameters["filter_rate"],
190190
}
191+
elif parameters["case_type"] == "LabelFilterPerformanceCase":
192+
custom_case_config = {
193+
"dataset_with_size_type": parameters["dataset_with_size_type"],
194+
"label_percentage": parameters["label_percentage"],
195+
}
191196
return custom_case_config
192197

193198

@@ -425,9 +430,9 @@ class CommonTypedDict(TypedDict):
425430
str,
426431
click.option(
427432
"--dataset-with-size-type",
428-
help="Dataset with size type for NewIntFilterPerformanceCase, you can use Medium Cohere (768dim, 1M)|"
429-
"Large Cohere (768dim, 10M)|Medium Bioasq (1024dim, 1M)|Large Bioasq (1024dim, 10M)|"
430-
"Large OpenAI (1536dim, 5M)|Medium OpenAI (1536dim, 500K)",
433+
help="Dataset with size type for NewIntFilterPerformanceCase/LabelFilterPerformanceCase, you can use "
434+
"Medium Cohere (768dim, 1M)|Large Cohere (768dim, 10M)|Medium Bioasq (1024dim, 1M)|"
435+
"Large Bioasq (1024dim, 10M)|Large OpenAI (1536dim, 5M)|Medium OpenAI (1536dim, 500K)",
431436
default="Medium Cohere (768dim, 1M)",
432437
show_default=True,
433438
),
@@ -441,6 +446,15 @@ class CommonTypedDict(TypedDict):
441446
show_default=True,
442447
),
443448
]
449+
label_percentage: Annotated[
450+
float,
451+
click.option(
452+
"--label-percentage",
453+
help="Filter rate for LabelFilterPerformanceCase",
454+
default=0.01,
455+
show_default=True,
456+
),
457+
]
444458

445459

446460
class HNSWBaseTypedDict(TypedDict):

0 commit comments

Comments
 (0)